CPU

model.py

@@ -156,7 +156,8 @@ class FoundModel(nn.Module):
     def decoder_load_weights(self, weights_path):
         print(f"Loading model from weights {weights_path}.")
         # Load states
-        state_dict = torch.load(weights_path)
+        map_location=torch.device('cpu') 
+        state_dict = torch.load(weights_path, map_location=map_location)
 
         # Decoder
         self.decoder.load_state_dict(state_dict["decoder"])