Bugfix: properly pass kwargs to torch.jit.load from `mlfow.pytorch.…

…load_model` (#5163)

* pass kwargs to `torch.jit.load`

Signed-off-by: Jake Schmidt <jacob.schmidt@recursionpharma.com>

* pop pickle_module

Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>

Co-authored-by: harupy <17039389+harupy@users.noreply.github.com>

mlflow/pytorch/__init__.py

@@ -635,7 +635,8 @@ 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)
+            kwargs.pop("pickle_module", None)  # `torch.jit.load` does not accept `pickle_module`.
+            return torch.jit.load(model_path, **kwargs)
 
 
 def load_model(model_uri, dst_path=None, **kwargs):