fix converter

InternLM · Dec 11, 2023 · f43c374 · f43c374
1 parent 9cd7044
commit f43c374
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/lmdeploy/turbomind/deploy/converter.py b/lmdeploy/turbomind/deploy/converter.py
@@ -117,7 +117,7 @@ def copy_tokenizer(model_path: str, tokenizer_path: str,
 def update_output_format(model_name: str, model_format: str, model_path: str,
                          output_format: str):
     """Update output format according to model info."""
-    TORCH_DTYPE_MAP = {'bfloat16': 'bf16'}
+    TORCH_DTYPE_MAP = {torch.bfloat16: 'bf16'}
     MODEL_NAME_MAP = {'qwen': 'bf16'}
     model_name = model_name.split('-')[0]