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: