fix import

Signed-off-by: Jimmy Zhang <jiemingz@nvidia.com>
JimmyZhang12 · Apr 12, 2024 · 5dde3cd · 5dde3cd
1 parent 7510089
commit 5dde3cd
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