From cf18e30316195e525dcb2dafba37025482fdbbd8 Mon Sep 17 00:00:00 2001 From: lanluo-nvidia Date: Tue, 5 Nov 2024 12:20:21 -0800 Subject: [PATCH] fix MutableTorchTensorRTModule load issue --- 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 )