small fixes

Author: zewenli98

py/torch_tensorrt/dynamo/_engine_caching.py

@@ -8,7 +8,7 @@
 
 import torch
 from torch._inductor.codecache import FxGraphCachePickler
-from torch.fx.experimental.proxy_tensor import maybe_disable_fake_tensor_mode
+from torch.fx.experimental.proxy_tensor import unset_fake_temporarily
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -34,7 +34,7 @@ def get_hash(gm: torch.fx.GraphModule) -> str:
             str: hash value of the GraphModule
         """
         # parameters are set to 0
-        with maybe_disable_fake_tensor_mode():
+        with unset_fake_temporarily():
             new_gm = copy.deepcopy(gm)
             for name, param in new_gm.named_parameters():
                 param.data.zero_()

py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py

@@ -6,6 +6,7 @@
 from typing import Any, Callable, Dict, List, NamedTuple, Optional, Sequence, Set, Tuple
 
 import numpy as np
+import tensorrt as trt
 import torch
 import torch.fx
 from torch.fx.node import _get_qualified_name
@@ -30,7 +31,6 @@
 from torch_tensorrt.fx.observer import Observer
 from torch_tensorrt.logging import TRT_LOGGER
 
-import tensorrt as trt
 from packaging import version
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
@@ -499,7 +499,7 @@ def run(
                     self._output_names = output_names
                     self.weight_name_map = weight_name_map
                     _LOGGER.info(
-                        "Hit the cached TRT engine. It is loaded and skip recompilation."
+                        "Found the cached engine that corresponds to this graph. It is directly loaded."
                     )
                     # TODO: refit the engine here or outside (within convert_module)?
                     return TRTInterpreterResult(