convert to float32 the loss if it is a bfloat16

microsoft · Nov 10, 2024 · 4e31c87 · 4e31c87
1 parent df1112f
commit 4e31c87
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/mttl/evaluators/loglike_evaluator.py b/mttl/evaluators/loglike_evaluator.py
@@ -68,7 +68,11 @@ def evaluate(
                 loss_per_option = compute_loglike_loss(
                     logits, batch["labels"], reduction="none"
                 )
-                loss_per_option = loss_per_option.cpu().numpy()
+                loss_per_option = loss_per_option.cpu()
+
+                if loss_per_option.dtype in [torch.bfloat16, torch.float16]:
+                    loss_per_option = loss_per_option.float().numpy()
+
                 loss_per_example = [
                     loss_per_option[
                         int(np.sum(num_options[:i])) : int(np.sum(num_options[: i + 1]))