hpcai-tech
/

OpenSora-STDiT-v1-HQ-16x256x256

Feature Extraction

Model card Files Files and versions Community

frankleeeee commited on Mar 22

Commit

8ef4ac6

•

1 Parent(s): 8d0ae18

Upload STDiT

Files changed (1) hide show

modeling_stdit.py +2 -1

modeling_stdit.py CHANGED Viewed

@@ -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)
         if self.enable_sequence_parallelism:
             x = gather_forward_split_backward(x, get_sequence_parallel_group(), dim=1, grad_scale="up")

                     tpe = self.pos_embed_temporal
             else:
                 tpe = None
+            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")