fix import

Signed-off-by: Jimmy Zhang <jiemingz@nvidia.com>
JimmyZhang12 · Apr 16, 2024 · 735d705 · 735d705
1 parent 10d8318
commit 735d705
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/export/tensorrt_llm.py b/nemo/export/tensorrt_llm.py
@@ -31,7 +31,7 @@
 from nemo.export.trt_llm.nemo.nemo_ckpt_convert import build_tokenizer
 from nemo.export.trt_llm.nemo_utils import get_tokenzier, nemo_llm_model_to_model_config, nemo_llm_to_model_config
 from nemo.export.trt_llm.tensorrt_llm_build import build_and_save_engine
-from nemo.export.trt_llm.tensorrt_llm_run import generate, generate_streaming, load, load_refit
+from nemo.export.trt_llm.tensorrt_llm_run import generate, generate_streaming, load, load_refit, create_gpt_session
 from nemo.export.trt_llm.utils import is_nemo_file, unpack_nemo_ckpt
 
 use_deploy = True