_is_bf16_available judgment supports npu
This commit is contained in:
parent
625a0e32c4
commit
384ab8db84
|
@ -37,7 +37,7 @@ from .logging import get_logger
|
||||||
|
|
||||||
_is_fp16_available = is_torch_npu_available() or is_torch_cuda_available()
|
_is_fp16_available = is_torch_npu_available() or is_torch_cuda_available()
|
||||||
try:
|
try:
|
||||||
_is_bf16_available = is_torch_bf16_gpu_available()
|
_is_bf16_available = is_torch_bf16_gpu_available() or (is_torch_npu_available() and torch.npu.is_bf16_supported())
|
||||||
except Exception:
|
except Exception:
|
||||||
_is_bf16_available = False
|
_is_bf16_available = False
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue