feat(naive_amp.py): not use fp32 output

yingtongxiong · Dec 25, 2023 · 053a13c · 053a13c
1 parent ec262ce
commit 053a13c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/internlm/core/naive_amp.py b/internlm/core/naive_amp.py
@@ -81,7 +81,7 @@ def _convert_to_fp16(self, input_: Any):
 
     def _convert_to_fp32(self, input_: Any):
         """Converts the input to fp32 if it is a Tensor of dtype float16."""
-        if isinstance(input_, Tensor) and input_.dtype in (torch.float16, torch.bfloat16):
+        if isinstance(input_, Tensor) and input_.dtype in (torch.float16,):
             input_ = input_.float()
         return input_