fix MutableTorchTensorRTModule load issue (#3281)

lanluo-nvidia · web-flow · commit 512968876973 · 2024-11-08T12:12:17.000-07:00
diff --git 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
         )

Original file line number	Diff line number	Diff line change
`@@ -498,7 +498,7 @@ def save(module: Any, path: str) -> None:`
`498`	`498`	`def load(path: str) -> Any:`
`499`	`499`	`# When the model get saved, init_finished is set to False.`
`500`	`500`	`# Class is restored to MutableTorchTensorRTModule, and some attribute is deleted`
`501`		`- module = torch.load(path)`
	`501`	`+ module = torch.load(path, weights_only=False)`
`502`	`502`	`module.pytorch_model = _make_refit_change_trigger(`
`503`	`503`	`module.original_model, module.refit_state`
`504`	`504`	`)`