load only weights

NVE · Aug 8, 2024 · f5fc23a · f5fc23a
1 parent 3f5b44c
commit f5fc23a
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/imagery/i.pytorch/pytorchlib/utils.py b/src/imagery/i.pytorch/pytorchlib/utils.py
@@ -149,12 +149,14 @@ def load_model(dl_model_path, dl_backbone, dl_kwargs, device="gpu"):
             dl_model.load_state_dict(
                 torch.load(
                     str(dl_model_path),
+                    weights_only=True,
                 )
             )
         else:
             dl_model.load_state_dict(
                 torch.load(
                     str(dl_model_path),
+                    weights_only=True,
                     map_location=torch.device("cpu"),
                 )
             )