Fixed issue where a buffer was stuck on cpu when offloading ideogram4

2026-06-07 13:31:03 -06:00
parent 1dc9a797cf
commit 35588726de
2 changed files with 5 additions and 1 deletions
--- a/extensions_built_in/diffusion_models/ideogram4/src/transformer.py
+++ b/extensions_built_in/diffusion_models/ideogram4/src/transformer.py
@@ -93,6 +93,10 @@ class Ideogram4MRoPE(nn.Module):
        # position_ids: (B, L, 3) of int.
        assert position_ids.ndim == 3 and position_ids.shape[-1] == 3
        batch_size, seq_len, _ = position_ids.shape
+        
+        if self.inv_freq.device == torch.device("cpu"):
+            # sometimes it gets stuck on CPU
+            self.inv_freq = self.inv_freq.to(position_ids.device)

        # (3, B, inv_freq_size, L)
        pos = position_ids.permute(2, 0, 1).to(dtype=torch.float32)