From b708b9175fbe9cd9eb32488b44a361cb72a86172 Mon Sep 17 00:00:00 2001
From: li-lizhe <147392333@qq.com>
Date: Sat, 5 Sep 2026 09:26:01 +0800
Subject: [PATCH] Fix Cosmos AdaLayerNorm autocast to follow the tensor device

CosmosAdaLayerNorm and CosmosAdaLayerNormZero wrapped their LayerNorm
call in torch.autocast(device_type="cuda", enabled=False). That context
only disables autocast for CUDA tensors; on other accelerators (Ascend
NPU, MPS, XPU, ROCm) it does not apply, so the norm can be computed in
the wrong precision and downstream dtype mismatches may appear.

Use hidden_states.device.type so the autocast context follows the actual
device. Device-agnostic and verified on NPU.

Fixes #1816
---
 fastvideo/models/dits/cosmos.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/fastvideo/models/dits/cosmos.py b/fastvideo/models/dits/cosmos.py
index d1cf64a4d8..1ef62ba8a1 100644
--- a/fastvideo/models/dits/cosmos.py
+++ b/fastvideo/models/dits/cosmos.py
@@ -91,7 +91,7 @@ def forward(self,
             embedded_timestep = embedded_timestep + temb[..., :2 * self.embedding_dim]
 
         shift, scale = embedded_timestep.chunk(2, dim=-1)
-        with torch.autocast(device_type="cuda", enabled=False):
+        with torch.autocast(device_type=hidden_states.device.type, enabled=False):
             hidden_states = self.norm(hidden_states)
 
         if embedded_timestep.ndim == 2:
@@ -131,7 +131,7 @@ def forward(
 
         shift, scale, gate = embedded_timestep.chunk(3, dim=-1)
 
-        with torch.autocast(device_type="cuda", enabled=False):
+        with torch.autocast(device_type=hidden_states.device.type, enabled=False):
             hidden_states = self.norm(hidden_states)
 
         if embedded_timestep.ndim == 2: