set weights_only=True to resolve security concern (#1409)

Closes BIO-15

---------

Signed-off-by: Peter St. John <pstjohn@nvidia.com>

Author: pstjohn

bionemo-recipes/recipes/codonfm_ptl_te/requirements.txt

@@ -156,7 +156,7 @@ notebook_shim==0.2.4
 notebook==7.4.5
 numpy==1.26.4
 nvidia-dali-cuda130==1.51.2
-nvidia-resiliency-ext==0.4.1+cuda13
+nvidia-resiliency-ext==0.4.1
 omegaconf==2.3.0
 opt_einsum==3.4.0
 optree==0.17.0

bionemo-recipes/recipes/codonfm_ptl_te/src/utils/load_checkpoint.py

@@ -36,7 +36,7 @@ def load_checkpoint(checkpoint_path: str, map_location: str = "cpu"):
             hparams = json.load(f)
         return {"state_dict": state_dict, "hyper_parameters": hparams}
     else:
-        return torch.load(checkpoint_path, map_location=map_location, weights_only=False)
+        return torch.load(checkpoint_path, map_location=map_location, weights_only=True)
 
 
 def download_checkpoint(repo_id: str, local_dir: str = None) -> str:  # noqa: RUF013

bionemo-recipes/recipes/codonfm_ptl_te/tests/utils/test_checkpoint.py

@@ -53,7 +53,7 @@ def test_load_pytorch_checkpoint_with_map_location(tmp_path):
         mock_torch_load.return_value = checkpoint_data
         load_checkpoint(str(checkpoint_path), map_location="cuda:0")
 
-        mock_torch_load.assert_called_once_with(str(checkpoint_path), map_location="cuda:0", weights_only=False)
+        mock_torch_load.assert_called_once_with(str(checkpoint_path), map_location="cuda:0", weights_only=True)
 
 
 def test_load_safetensors_directory(tmp_path):

bionemo-recipes/recipes/esm2_native_te/checkpoint.py

@@ -124,7 +124,7 @@ def load_checkpoint_ddp(
     checkpoint = torch.load(
         checkpoint_path / "checkpoint.pt",
         map_location=f"cuda:{dist_config.local_rank}",
-        weights_only=False,
+        weights_only=True,
     )
 
     model.load_state_dict(checkpoint["model"])

bionemo-recipes/recipes/llama3_native_te/checkpoint.py

@@ -124,7 +124,7 @@ def load_checkpoint_ddp(
     checkpoint = torch.load(
         checkpoint_path / "checkpoint.pt",
         map_location=f"cuda:{dist_config.local_rank}",
-        weights_only=False,
+        weights_only=True,
     )
 
     model.load_state_dict(checkpoint["model"])

bionemo-recipes/recipes/vit/checkpoint.py

@@ -36,7 +36,7 @@ def load_torch_checkpoint(checkpoint_path, model, megatron_fsdp=False):
     """
     # Load model checkpoint. Must load with weights_only=False
     # if you have an optimizer state in your checkpoint.
-    checkpoint = torch.load(checkpoint_path, weights_only=False)
+    checkpoint = torch.load(checkpoint_path, weights_only=True)
     # Remove the "module." prefix from the keys of checkpoints
     # derived from Megatron-FSDP.
     # TODO(@cspades): Remove this when the Megatron-FSDP checkpoint naming is fixed.