fix torch warning

segasai · Jul 30, 2024 · ea3f3de · ea3f3de
1 parent aaa649f
commit ea3f3de
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/py/rvspecfit/nn/RVSInterpolator.py b/py/rvspecfit/nn/RVSInterpolator.py
@@ -22,7 +22,8 @@ def __init__(self, kwargs):
             self.device = torch.device(device_env)
         self.nni.load_state_dict(
             torch.load(kwargs['template_lib'] + '/' + kwargs['nn_file'],
-                       map_location=self.device))
+                       map_location=self.device,
+                       weights_only=True))
         # self.device = list(self.nni.children())[0][0].pc_layer.weight.device
 
         self.nni.to(self.device)