fix typehint

Author: Mark Obozov

src/forge/losses/grpo_loss.py

@@ -16,7 +16,7 @@ class SimpleGRPOLoss(nn.Module):
         https://github.com/huggingface/trl/blob/417915a3e4d3e3bc8d7b196594308b8eabf928be/trl/trainer/grpo_trainer.py#L1624.
     """
 
-    def __init__(self, beta: float = 0.1) -> torch.Tensor | LossMetrics:
+    def __init__(self, beta: float = 0.1) -> tuple[torch.Tensor, LossMetrics]:
         super().__init__()
         self.beta = beta
 

src/forge/losses/reinforce_loss.py

@@ -30,7 +30,7 @@ def __init__(self):
 
     def forward(
         self, trainer_logits, target_ids, target_mask, target_weights, target_log_probs
-    ) -> torch.Tensor | LossMetrics:
+    ) -> tuple[torch.Tensor, LossMetrics]:
         trainer_log_probs = selective_log_softmax(trainer_logits, target_ids)
         target_mask = target_mask.detach()
         target_weights = target_weights