fix Ruff, lint

Author: Huy Vu2

dfm/src/megatron/data/common/base_energon_datamodule.py

@@ -336,4 +336,4 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
         update_num_microbatches(
             consumed_samples=consumed_samples,
             consistency_check=False,
-        )
+        )

dfm/src/megatron/data/common/diffusion_energon_datamodule.py

@@ -182,4 +182,4 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
         try:
             super().load_state_dict(state_dict)
         except Exception as e:
-            logging.warning(f"datamodule.load_state_dict failed  {e}")
+            logging.warning(f"datamodule.load_state_dict failed  {e}")

dfm/src/megatron/data/common/diffusion_sample.py

@@ -93,4 +93,4 @@ def __lt__(self, other: Any) -> bool:
             return self.seq_len_q.item() < other.seq_len_q.item()
         elif isinstance(other, int):
             return self.seq_len_q.item() < other
-        raise NotImplementedError
+        raise NotImplementedError

dfm/src/megatron/data/common/diffusion_task_encoder_with_sp.py

@@ -109,4 +109,4 @@ def cat(attr):
 
     @stateless
     def batch(self, samples: List[DiffusionSample]) -> dict:
-        raise NotImplementedError
+        raise NotImplementedError

dfm/src/megatron/data/common/sequence_packing_utils.py

@@ -102,4 +102,4 @@ def concat_pad(tensor_list, max_seq_length):
         result[current_index : current_index + length] = tensor
         current_index += length
 
-    return result
+    return result

dfm/src/megatron/data/wan/wan_taskencoder.py

@@ -22,7 +22,7 @@
 from megatron.core import parallel_state
 from megatron.energon import SkipSample
 from megatron.energon.task_encoder.base import stateless
-from megatron.energon.task_encoder.cooking import Cooker, basic_sample_keys
+from megatron.energon.task_encoder.cooking import basic_sample_keys, Cooker
 import torch
 import torch.nn.functional as F
 

dfm/src/megatron/model/wan/rope_utils.py

@@ -86,4 +86,4 @@ def forward(self, n_head, dim_head, cu_seqlens_q_padded, grid_sizes, device):
         # we don't need to scatter the freqs to the context parallel region,
         # because mcore rope_utils will automatically retrieve the correct freqs for each context parallel region
 
-        return freqs_real
+        return freqs_real

examples/megatron/recipes/wan/prepare_energon_dataset_wan.py

@@ -18,11 +18,11 @@
 from typing import Dict, List, Optional, Tuple
 
 import cv2
-from diffusers import AutoencoderKLWan
 import numpy as np
 import torch
-from transformers import AutoTokenizer, UMT5EncoderModel
 import webdataset as wds
+from diffusers import AutoencoderKLWan
+from transformers import AutoTokenizer, UMT5EncoderModel
 
 
 def _map_interpolation(resize_mode: str) -> int: