lint

Signed-off-by: ashors1 <ashors@nvidia.com>

Author: ashors1

nemo_rl/models/megatron/data.py

@@ -58,6 +58,7 @@ class ProcessedMicrobatch:
         packed_seq_params: PackedSeqParams for sequence packing (None if not packing)
         cu_seqlens_padded: Padded cumulative sequence lengths (None if not packing)
     """
+
     data_dict: BatchedDataDict[Any]
     input_ids: torch.Tensor
     input_ids_cp_sharded: torch.Tensor
@@ -202,6 +203,7 @@ def get_microbatch_iterator(
         padded_seq_length,
     )
 
+
 def process_microbatch(
     data_dict: BatchedDataDict[Any],
     seq_length_key: Optional[str] = None,
@@ -338,6 +340,7 @@ def process_global_batch(
     }
 
 
+
 def _pack_sequences_for_megatron(
     input_ids: torch.Tensor,
     seq_lengths: torch.Tensor,

nemo_rl/models/policy/workers/megatron_policy_worker.py

@@ -85,7 +85,6 @@
     LogprobsPostProcessor,
     TopkLogitsPostProcessor,
 )
-from nemo_rl.models.megatron.community_import import import_model_from_hf_name
 from nemo_rl.models.policy import PolicyConfig
 from nemo_rl.models.policy.interfaces import (
     ColocatablePolicyInterface,

tests/unit/models/megatron/test_megatron_data.py

@@ -1709,4 +1709,3 @@ def test_get_pack_sequence_parameters_for_megatron(get_pack_sequence_parameters_
     # Check that all workers succeeded
     for i, result in enumerate(results):
         assert result["success"], f"Worker {i} failed: {result['error']}"
-