From 512968876973c7bfbd30bd38750e0ade9f54b677 Mon Sep 17 00:00:00 2001 From: lanluo-nvidia Date: Fri, 8 Nov 2024 11:12:17 -0800 Subject: [PATCH] fix MutableTorchTensorRTModule load issue (#3281) --- py/torch_tensorrt/dynamo/runtime/_MutableTorchTensorRTModule.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 )