added map_location=cpu

ViCCo-Group · Sep 10, 2024 · 9f8a6ce · 9f8a6ce
1 parent cd80049
commit 9f8a6ce
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/thingsvision/core/extraction/extractors.py b/thingsvision/core/extraction/extractors.py
@@ -470,6 +470,7 @@ def load_model_from_source(self) -> None:
                 # load and cache state_dict
                 state_dict = torch.hub.load_state_dict_from_url(
                     model_config["checkpoint_url"],
+                    map_location=torch.device("cpu"),
                     # IMPORTANT that this is unique as it will be used for caching
                     file_name=unique_model_filename
                 )