torchfix lint

1 year ago · 97f1a52f31
parent fd28876f6d
commit 97f1a52f31
1 changed files with 2 additions and 1 deletions
--- a/swarms/utils/load_model_torch.py
+++ b/swarms/utils/load_model_torch.py
@ -38,13 +38,14 @@ def load_model_torch(
    try:
        if model is None:
            model = torch.load(
-                model_path, map_location=map_location, *args, **kwargs
+                model_path, map_location=map_location, weights_only=True, *args, **kwargs
            )
        else:
            model.load_state_dict(
                torch.load(
                    model_path,
                    map_location=map_location,
+                    weights_only=True,
                    *args,
                    **kwargs,
                ),