image classification - fix bug in zoo stub load (#959)

bfineran · web-flow · commit a9d0585d823b · 2022-07-29T17:28:38.000-04:00
diff --git a/src/sparseml/pytorch/image_classification/utils/helpers.py b/src/sparseml/pytorch/image_classification/utils/helpers.py
@@ -271,7 +271,7 @@ def create_model(
     """
     with torch_distributed_zero_first(local_rank):
         # only download once locally
-        if checkpoint_path and checkpoint_path.lower() == "zoo":
+        if checkpoint_path and checkpoint_path.lower().startswith("zoo"):
             checkpoint_path = _download_model_from_zoo_using_recipe(
                 recipe_stub=recipe_path,
             )

Original file line number	Diff line number	Diff line change
`@@ -271,7 +271,7 @@ def create_model(`
`271`	`271`	`"""`
`272`	`272`	`with torch_distributed_zero_first(local_rank):`
`273`	`273`	`# only download once locally`
`274`		`- if checkpoint_path and checkpoint_path.lower() == "zoo":`
	`274`	`+ if checkpoint_path and checkpoint_path.lower().startswith("zoo"):`
`275`	`275`	`checkpoint_path = _download_model_from_zoo_using_recipe(`
`276`	`276`	`recipe_stub=recipe_path,`
`277`	`277`	`)`