Fix torch.compile on transfomer

iddl · iddl · commit d62ee1ec7307 · 2025-09-26T10:50:05.000-04:00
diff --git a/src/diffusers/models/transformers/transformer_chroma.py b/src/diffusers/models/transformers/transformer_chroma.py
@@ -169,8 +169,8 @@ def forward(self, timestep: torch.Tensor) -> torch.Tensor:
         batch_size = timestep.shape[0]
 
         timesteps_proj = self.time_proj(timestep).to(dtype=timestep.dtype)
-        guidance_proj = self.guidance_proj(torch.tensor([0] * batch_size)).to(
-            dtype=timestep.dtype, device=timestep.device
+        guidance_proj = self.guidance_proj(torch.tensor([0] * batch_size, device=timestep.device)).to(
+            dtype=timestep.dtype
         )
 
         mod_proj = self.mod_proj.to(dtype=timesteps_proj.dtype, device=timesteps_proj.device).repeat(batch_size, 1, 1)

Original file line number	Diff line number	Diff line change
`@@ -169,8 +169,8 @@ def forward(self, timestep: torch.Tensor) -> torch.Tensor:`
`169`	`169`	`batch_size = timestep.shape[0]`
`170`	`170`
`171`	`171`	`timesteps_proj = self.time_proj(timestep).to(dtype=timestep.dtype)`
`172`		`- guidance_proj = self.guidance_proj(torch.tensor([0] * batch_size)).to(`
`173`		`- dtype=timestep.dtype, device=timestep.device`
	`172`	`+ guidance_proj = self.guidance_proj(torch.tensor([0] * batch_size, device=timestep.device)).to(`
	`173`	`+ dtype=timestep.dtype`
`174`	`174`	`)`
`175`	`175`
`176`	`176`	`mod_proj = self.mod_proj.to(dtype=timesteps_proj.dtype, device=timesteps_proj.device).repeat(batch_size, 1, 1)`