Update model.py

model.py

@@ -167,7 +167,7 @@ class MultiHeadClassification(nn.Module):
         Returns:
             None
         """
-        model = torch.load(path)
+        model = torch.load(path, map_location=self.device)
         if head_name in self.heads:
             num_classes = model['weight'].shape[0]
             self.heads[head_name].load_state_dict(model)
@@ -209,7 +209,7 @@ class MultiHeadClassification(nn.Module):
         Args:
             path (str): Path to the file
         """
-        self.load_state_dict(torch.load(path))
+        self.load_state_dict(torch.load(path, map_location=self.device))
         self.to(self.torch_dtype).to(self.device)
 
     def save_backbone(self, path):