diff --git a/comfy/ldm/cosmos/blocks.py b/comfy/ldm/cosmos/blocks.py
index ff653c14f..a12f892d2 100644
--- a/comfy/ldm/cosmos/blocks.py
+++ b/comfy/ldm/cosmos/blocks.py
@@ -796,15 +796,12 @@ class GeneralDITTransformerBlock(nn.Module):
         adaln_lora_B_3D: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         for block in self.blocks:
-            if self.training:
-                x = torch.utils.checkpoint.checkpoint(block, x, emb_B_D, crossattn_emb, crossattn_mask, rope_emb_L_1_1_D, adaln_lora_B_3D, use_reentrant=False)
-            else:
-                x = block(
-                    x,
-                    emb_B_D,
-                    crossattn_emb,
-                    crossattn_mask,
-                    rope_emb_L_1_1_D=rope_emb_L_1_1_D,
-                    adaln_lora_B_3D=adaln_lora_B_3D,
-                )
+            x = block(
+                x,
+                emb_B_D,
+                crossattn_emb,
+                crossattn_mask,
+                rope_emb_L_1_1_D=rope_emb_L_1_1_D,
+                adaln_lora_B_3D=adaln_lora_B_3D,
+            )
         return x