fix some warnings

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -5270,7 +5270,7 @@ def load(
 
         assert path.exists() and path.is_file()
 
-        package = torch.load(str(path), map_location = map_location)
+        package = torch.load(str(path), map_location = map_location, weights_only = True)
 
         model_package = package['model']
         current_version = version('alphafold3_pytorch')
@@ -5293,7 +5293,7 @@ def init_and_load(
 
         assert path.is_file()
 
-        package = torch.load(str(path), map_location = map_location)
+        package = torch.load(str(path), map_location = map_location, weights_only = True)
 
         model_package = package['model']
 

alphafold3_pytorch/trainer.py

@@ -584,7 +584,7 @@ def load(
 
         # load optimizer and scheduler states
 
-        package = torch.load(str(path))
+        package = torch.load(str(path), weights_only = True)
 
         unwrapped_optimizer = _unwrap_objects(self.optimizer)