weights_only

hstern2 · Sep 20, 2024 · 5214092 · 5214092
1 parent ad6f12a
commit 5214092
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/amsr/lstm.py b/amsr/lstm.py
@@ -151,7 +151,7 @@ def train_and_save_model(
 
     @classmethod
     def from_saved_model(cls, model_path):
-        checkpoint = torch.load(model_path, map_location=_device())
+        checkpoint = torch.load(model_path, map_location=_device(), weights_only=True)
         model = cls(
             d_model=checkpoint["d_model"],
             nhid=checkpoint["nhid"],