_is_bf16_available judgment supports npu
Former-commit-id: 50a1e892a1005b4cdd82dca1005f71db08ed89a2
This commit is contained in:
@@ -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
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user