Skip to content

Commit

Permalink
fix MutableTorchTensorRTModule load issue (#3281)
Browse files Browse the repository at this point in the history
  • Loading branch information
lanluo-nvidia authored Nov 8, 2024
1 parent eb6d410 commit 5129688
Showing 1 changed file with 1 addition and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
Expand Down

0 comments on commit 5129688

Please sign in to comment.