Problem
and in wrap the LayerNorm in a hardcoded CUDA autocast context:
only disables autocast for CUDA tensors. On other accelerators (Ascend NPU, MPS, XPU, ROCm) the autocast context does not apply, so the norm may be computed in the wrong precision (e.g. bf16) instead of the intended fp32, which can cause numerical differences or dtype mismatches downstream. It also hardcodes a CUDA-specific assumption.
Expected
Disable autocast for the tensor’s actual device: . This is device-agnostic and works on NPU/MPS/XPU/CUDA alike (verified: is valid).
Environment
- repo: hao-ai-lab/FastVideo
- device: Ascend 910B (torch.accelerator.current_accelerator() == npu)
Problem
and in wrap the LayerNorm in a hardcoded CUDA autocast context:
only disables autocast for CUDA tensors. On other accelerators (Ascend NPU, MPS, XPU, ROCm) the autocast context does not apply, so the norm may be computed in the wrong precision (e.g. bf16) instead of the intended fp32, which can cause numerical differences or dtype mismatches downstream. It also hardcodes a CUDA-specific assumption.
Expected
Disable autocast for the tensor’s actual device: . This is device-agnostic and works on NPU/MPS/XPU/CUDA alike (verified: is valid).
Environment