Use weights_only for load (#127)

pytorch-labs · Sep 13, 2024 · e6aadeb · e6aadeb
1 parent b27b84d
commit e6aadeb
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/experiments/eval_combo.py b/experiments/eval_combo.py
@@ -365,7 +365,7 @@ def run(
         weights_path = Path(f"static_quant_scalars/{sam_model_type}_{batch_size}_static_quant_weights.ptk")
         if weights_path.exists() and weights_path.is_file():
             print("Loading static quantization weights")
-            weights = torch.load(f"static_quant_scalars/{sam_model_type}_{batch_size}_static_quant_weights.ptk")
+            weights = torch.load(f"static_quant_scalars/{sam_model_type}_{batch_size}_static_quant_weights.ptk", weights_only=True)
             from static_quant import set_x_absmax
             set_x_absmax(predictor.model.image_encoder, weights)
     elif compress == "sparse":

diff --git a/segment_anything_fast/build_sam.py b/segment_anything_fast/build_sam.py
@@ -142,6 +142,6 @@ def _build_sam(
     sam.eval()
     if checkpoint is not None:
         with open(checkpoint, "rb") as f:
-            state_dict = torch.load(f)
+            state_dict = torch.load(f, weights_only=True)
         sam.load_state_dict(state_dict)
     return sam