Improve the safety of torch.load with weights_only=True

hankcs · Oct 8, 2024 · 5008d7b · 5008d7b
1 parent 32428a2
commit 5008d7b
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/hanlp/common/torch_component.py b/hanlp/common/torch_component.py
@@ -97,7 +97,10 @@ def load_weights(self, save_dir, filename='model.pt', **kwargs):
         save_dir = get_resource(save_dir)
         filename = os.path.join(save_dir, filename)
         # flash(f'Loading model: {filename} [blink]...[/blink][/yellow]')
-        self.model_.load_state_dict(torch.load(filename, map_location='cpu'), strict=False)
+        try:
+            self.model_.load_state_dict(torch.load(filename, map_location='cpu', weights_only=True), strict=False)
+        except TypeError:
+            self.model_.load_state_dict(torch.load(filename, map_location='cpu'), strict=False)
         # flash('')
 
     def save_config(self, save_dir, filename='config.json'):