diff --git a/py/torch_tensorrt/dynamo/runtime/_MutableTorchTensorRTModule.py b/py/torch_tensorrt/dynamo/runtime/_MutableTorchTensorRTModule.py index 9abd896d50..28f0954185 100644 --- a/py/torch_tensorrt/dynamo/runtime/_MutableTorchTensorRTModule.py +++ b/py/torch_tensorrt/dynamo/runtime/_MutableTorchTensorRTModule.py @@ -498,7 +498,7 @@ def save(module: Any, path: str) -> None: def load(path: str) -> Any: # When the model get saved, init_finished is set to False. # Class is restored to MutableTorchTensorRTModule, and some attribute is deleted - module = torch.load(path) + module = torch.load(path, weights_only=False) module.pytorch_model = _make_refit_change_trigger( module.original_model, module.refit_state )