Specify map location when loading model

tests/conftest.py

@@ -45,7 +45,7 @@ def __init__(
             backbone = torch.nn.Linear(num_frames, num_hidden)
             torch.nn.init.ones_(backbone.weight)
         else:
-            backbone = self.load_from_checkpoint(finetune_from).backbone
+            backbone = self.load_from_checkpoint(finetune_from, map_location=self.device).backbone
 
         for param in backbone.parameters():
             param.requires_grad = False

zamba/models/efficientnet_models.py

@@ -17,15 +17,20 @@ class TimeDistributedEfficientNet(ZambaVideoClassificationLightningModule):
     )
 
     def __init__(
-        self, num_frames=16, finetune_from: Optional[Union[os.PathLike, str]] = None, **kwargs
+        self,
+        num_frames=16,
+        finetune_from: Optional[Union[os.PathLike, str]] = None,
+        **kwargs,
     ):
         super().__init__(**kwargs)
 
         if finetune_from is None:
             efficientnet = timm.create_model("efficientnetv2_rw_m", pretrained=True)
             efficientnet.classifier = nn.Identity()
         else:
-            efficientnet = self.load_from_checkpoint(finetune_from).base.module
+            efficientnet = self.load_from_checkpoint(
+                finetune_from, map_location=self.device
+            ).base.module
 
         # freeze base layers
         for param in efficientnet.parameters():

zamba/models/model_manager.py

@@ -62,7 +62,7 @@ def instantiate_model(
            Only used if labels is not None.
         model_name (ModelEnum, optional): Model name used to look up default hparams used for that model.
             Only relevant if training from scratch.
-        use_default_model_labels(bool, optional): Whether to output the full set of default model labels rather than
+        use_default_model_labels (bool, optional): Whether to output the full set of default model labels rather than
             just the species in the labels file. Only used if labels is not None.
 
     Returns:
@@ -110,10 +110,8 @@ def instantiate_model(
             return resume_training(
                 scheduler_config=scheduler_config,
                 hparams=hparams,
-                species=species,
                 model_class=model_class,
                 checkpoint=checkpoint,
-                labels=labels,
             )
 
         else:
@@ -157,10 +155,8 @@ def replace_head(scheduler_config, hparams, species, model_class, checkpoint):
 def resume_training(
     scheduler_config,
     hparams,
-    species,
     model_class,
     checkpoint,
-    labels,
 ):
     # resume training; add additional species columns to labels file if needed
     logger.info(

zamba/models/slowfast_models.py

@@ -55,7 +55,7 @@ def __init__(
         if finetune_from is None:
             self.initialize_from_torchub()
         else:
-            model = self.load_from_checkpoint(finetune_from)
+            model = self.load_from_checkpoint(finetune_from, map_location=self.device)
             self._backbone_output_dim = model.head.proj.in_features
             self.backbone = model.backbone
             self.base = model.base