Quick fix to typing issue (with TODO item for better fix)

odelalleau · odelalleau · commit 179767e35d54 · 2025-07-25T10:43:15.000-04:00
Signed-off-by: Olivier Delalleau &lt;507137+odelalleau@users.noreply.github.com&gt;
diff --git a/nemo_rl/algorithms/loss_functions.py b/nemo_rl/algorithms/loss_functions.py
@@ -623,7 +623,8 @@ def _dpo_loss(
             rewards, sample_mask, global_valid_seqs, self.reference_policy_kl_penalty
         )
 
-    def __call__(
+    # TODO a cleaner typing fix would be required (probably that DPOLossFn should not inherit from PreferenceLoss)
+    def __call__( # type: ignore
         self,
         next_token_logits: Tensor,
         data: BatchedDataDict[DPOLossDataDict],

Original file line number	Diff line number	Diff line change
`@@ -623,7 +623,8 @@ def _dpo_loss(`
`623`	`623`	`rewards, sample_mask, global_valid_seqs, self.reference_policy_kl_penalty`
`624`	`624`	`)`
`625`	`625`
`626`		`- def __call__(`
	`626`	`+ # TODO a cleaner typing fix would be required (probably that DPOLossFn should not inherit from PreferenceLoss)`
	`627`	`+ def __call__( # type: ignore`
`627`	`628`	`self,`
`628`	`629`	`next_token_logits: Tensor,`
`629`	`630`	`data: BatchedDataDict[DPOLossDataDict],`