[image classification] bug fix - use full zoo stub prefix to determine zoo checkpoint (#1036)

bfineran · web-flow · commit 3c042f361de6 · 2022-09-15T09:44:05.000-04:00
prior to this fix, IC trainer would recognize any checkpoint starting with "zoo" as a stub, this causes issues if a local path starts with zoo. this PR fixes the issue by matching to the full zoo prefix "zoo:"
diff --git a/src/sparseml/pytorch/image_classification/utils/trainer.py b/src/sparseml/pytorch/image_classification/utils/trainer.py
@@ -332,7 +332,7 @@ def _run_train_epoch(
         )
 
     def _setup_checkpoint_manager(self):
-        if self.checkpoint_path and self.checkpoint_path.startswith("zoo"):
+        if self.checkpoint_path and self.checkpoint_path.startswith("zoo:"):
             zoo_model = Model(self.checkpoint_path)
             self.checkpoint_path = download_framework_model_by_recipe_type(zoo_model)
 

Original file line number	Diff line number	Diff line change
`@@ -332,7 +332,7 @@ def _run_train_epoch(`
`332`	`332`	`)`
`333`	`333`
`334`	`334`	`def _setup_checkpoint_manager(self):`
`335`		`- if self.checkpoint_path and self.checkpoint_path.startswith("zoo"):`
	`335`	`+ if self.checkpoint_path and self.checkpoint_path.startswith("zoo:"):`
`336`	`336`	`zoo_model = Model(self.checkpoint_path)`
`337`	`337`	`self.checkpoint_path = download_framework_model_by_recipe_type(zoo_model)`
`338`	`338`