From e7a5766a771c62f0da0a36d6e6bc4588948096b1 Mon Sep 17 00:00:00 2001 From: Sergii Dymchenko Date: Thu, 15 Aug 2024 11:40:16 -0700 Subject: [PATCH] Use weights_only for load --- experiments/eval_combo.py | 2 +- segment_anything_fast/build_sam.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/experiments/eval_combo.py b/experiments/eval_combo.py index 91845c2..3fa08b3 100644 --- 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 index a9a4c99..f27c270 100644 --- 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