From fa55698269c70f7a2b48224392bc5fef0174f5b5 Mon Sep 17 00:00:00 2001 From: Waruna Wickramasingha Date: Tue, 17 Dec 2024 12:09:39 +0000 Subject: [PATCH] torch.load future warning fixed --- diffraction/WISH/bragg-detect/cnn/BraggDetectCNN.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diffraction/WISH/bragg-detect/cnn/BraggDetectCNN.py b/diffraction/WISH/bragg-detect/cnn/BraggDetectCNN.py index 6d9a64d..83b7794 100644 --- a/diffraction/WISH/bragg-detect/cnn/BraggDetectCNN.py +++ b/diffraction/WISH/bragg-detect/cnn/BraggDetectCNN.py @@ -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)