Upload STDiT

modeling_stdit.py

@@ -148,7 +148,8 @@ class STDiT(PreTrainedModel):
                     tpe = self.pos_embed_temporal
             else:
                 tpe = None
-            x = auto_grad_checkpoint(block, x, y, t0, y_lens, tpe)
+            x = block(x, y, t0, y_lens, tpe)
+            # x = auto_grad_checkpoint(block, x, y, t0, y_lens, tpe)
 
         if self.enable_sequence_parallelism:
             x = gather_forward_split_backward(x, get_sequence_parallel_group(), dim=1, grad_scale="up")