Skip to content

Commit

Permalink
fix MutableTorchTensorRTModule load issue
Browse files Browse the repository at this point in the history
  • Loading branch information
lanluo-nvidia committed Nov 5, 2024
1 parent 8e2c82d commit cf18e30
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 cf18e30

Please sign in to comment.