Try autodetecting dtype parameter by default

Signed-off-by: Jan Lasek <[email protected]>
NVIDIA · Jan 21, 2025 · 1ad05bc · 1ad05bc
1 parent c84be55
commit 1ad05bc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/export/tensorrt_llm.py b/nemo/export/tensorrt_llm.py
@@ -171,7 +171,7 @@ def export(
         paged_kv_cache: bool = True,
         remove_input_padding: bool = True,
         paged_context_fmha: bool = False,
-        dtype: Optional[str] = "bfloat16",
+        dtype: Optional[str] = None,
         load_model: bool = True,
         use_lora_plugin: str = None,
         lora_target_modules: List[str] = None,