QWEN: Fix unsupported ScalarType BFloat16 (ml-explore#187)

Fix unsupported ScalarType BFloat16.
Blaizzy · Dec 25, 2023 · 447abb2 · 447abb2
1 parent c190af6
commit 447abb2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/llms/qwen/convert.py b/llms/qwen/convert.py
@@ -60,7 +60,7 @@ def convert(args):
         args.model, trust_remote_code=True, torch_dtype=torch.float16
     )
     state_dict = model.state_dict()
-    weights = {replace_key(k): v.numpy() for k, v in state_dict.items()}
+    weights = {replace_key(k): (v.numpy() if v.dtype != torch.bfloat16 else v.to(torch.float32).numpy()) for k, v in state_dict.items()}
     config = model.config.to_dict()
 
     if args.quantize: