pass kwargs to torch.jit.load

Signed-off-by: Jake Schmidt <jacob.schmidt@recursionpharma.com>
mlflow · Dec 14, 2021 · 987298f · 987298f
1 parent 0dbd5c4
commit 987298f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mlflow/pytorch/__init__.py b/mlflow/pytorch/__init__.py
@@ -635,7 +635,7 @@ def _load_model(path, **kwargs):
             return torch.load(model_path, **kwargs)
         except Exception:
             # If fails, assume the model as a scripted model
-            return torch.jit.load(model_path)
+            return torch.jit.load(model_path, **kwargs)
 
 
 def load_model(model_uri, dst_path=None, **kwargs):