Skip to content

Commit

Permalink
torch.load future warning fixed
Browse files Browse the repository at this point in the history
  • Loading branch information
warunawickramasingha committed Jan 2, 2025
1 parent 4a2ba29 commit fa55698
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion diffraction/WISH/bragg-detect/cnn/BraggDetectCNN.py
Original file line number Diff line number Diff line change
Expand Up @@ -168,7 +168,7 @@ def _select_device(self):

def _load_cnn_model_from_weights(self, weights_path):
model = self._get_fasterrcnn_resnet50_fpn(num_classes=2)
model.load_state_dict(tc.load(weights_path, map_location=self.device))
model.load_state_dict(tc.load(weights_path, map_location=self.device, weights_only=True))
return model.to(self.device)


Expand Down

0 comments on commit fa55698

Please sign in to comment.