fix fsdp to cpu

Author: Edenzzzz

fastvideo/v1/layers/layernorm.py

@@ -74,7 +74,7 @@ def forward_native(
             # TODO(wenxuan): When using CPU offload, FSDP has a bug that doesn't unwrap DTensor in final_layer_norm.
             # Report this
             if isinstance(self.weight, DTensor):
-                x = x * self.weight.to(x.device).full_tensor()
+                x = x * self.weight.to_local().to(x.device)
             else:
                 x = x * self.weight
         if residual is None:

fastvideo/v1/pipelines/stages/text_encoding.py

@@ -107,9 +107,6 @@ def forward(
                     batch.negative_attention_mask.append(
                         negative_attention_mask)
 
-            if fastvideo_args.text_encoder_offload:
-                text_encoder.to('cpu')
-
         return batch
 
     def verify_input(self, batch: ForwardBatch,