Merge pull request #12 from aritraghsh09/Autocrop-no-gpu

Updating autocrop.py so that it can run on CPUs as well
aritraghsh09 · May 16, 2024 · 5409d31 · 5409d31
2 parents 0d3f24a + e6b81b7
commit 5409d31
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/ggt/modules/autocrop.py b/ggt/modules/autocrop.py
@@ -90,7 +90,8 @@ def main(
     model = model.to(device)
 
     # Load the model from a saved state if provided
-    model.load_state_dict(torch.load(model_path))
+    model.load_state_dict(torch.load(model_path,
+                                     map_location=torch.device(device)))
 
     # Collect all images, then iterate
     images = glob.glob(str(Path(image_dir) / "*.fits"))