Add weights_only param

SamSamhuns · Aug 7, 2024 · b489b62 · b489b62
1 parent bbb4e09
commit b489b62
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/face_detection_and_extraction/modules/blazeface/blazeface.py b/face_detection_and_extraction/modules/blazeface/blazeface.py
@@ -232,7 +232,7 @@ def _device(self):
         return self.classifier_8.weight.device
 
     def load_weights(self, path):
-        self.load_state_dict(torch.load(path))
+        self.load_state_dict(torch.load(path, weights_only=True))
         self.eval()
 
     def load_anchors(self, path, use_numpy=False):

diff --git a/face_detection_and_extraction/modules/yolov5_face/pytorch/models/experimental.py b/face_detection_and_extraction/modules/yolov5_face/pytorch/models/experimental.py
@@ -118,7 +118,7 @@ def attempt_load(weights, map_location=None):
     # Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a
     model = Ensemble()
     for w in weights if isinstance(weights, list) else [weights]:
-        model.append(torch.load(w, map_location=map_location)[
+        model.append(torch.load(w, map_location=map_location, weights_only=False)[
                      'model'].float().fuse().eval())  # load FP32 model
 
     # Compatibility updates