未验证 提交 60d7b0a3 编写于 作者: C CurryRice233 提交者: GitHub

add npu support dtypes (#4223)

Co-authored-by: Njializheng <jializheng@huawei.com>
上级 42c1e916
......@@ -134,6 +134,9 @@ class NPU_Accelerator(DeepSpeedAccelerator):
def is_fp16_supported(self):
return True
def supported_dtypes(self):
return [torch.float, torch.half, torch.bfloat16]
# Misc
def amp(self):
if hasattr(torch.npu, 'amp'):
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册