Fixup model extraction code.

vanvalenlab · Apr 5, 2024 · a84bd37 · a84bd37
1 parent 110e01e
commit a84bd37
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/cellSAM/model.py b/cellSAM/model.py
@@ -44,7 +44,9 @@ def get_model(model: nn.Module = None) -> nn.Module:
             cellsam_assets_dir.mkdir(parents=True, exist_ok=True)
         if not model_path.exists():
             fetch_data("models/cellsam_base.tar.gz", cache_subdir="models")
-            extract_archive(model_path, cellsam_assets_dir)
+            extract_archive(
+                cellsam_assets_dir / "cellsam_base.tar.gz", cellsam_assets_dir
+            )
             assert model_path.exists()
         model = CellSAM(config)
     model.load_state_dict(torch.load(model_path))