From 244e1b5ba331cb4c1ed96d88d0895c252567f7f3 Mon Sep 17 00:00:00 2001 From: Sylvain Gugger <35901082+sgugger@users.noreply.github.com> Date: Tue, 22 Sep 2020 09:20:03 -0400 Subject: [PATCH] Fix #7304 (#7305) --- src/transformers/trainer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index bb17635171..a58ac0a427 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -1334,9 +1334,9 @@ class Trainer: elif is_torch_tpu_available(): # tpu-comment: Get all predictions and labels from all worker shards of eval dataset if preds is not None: - preds = nested_xla_mesh_reduce("eval_preds", preds) + preds = nested_xla_mesh_reduce(preds, "eval_preds") if label_ids is not None: - label_ids = nested_xla_mesh_reduce("eval_label_ids", label_ids, torch.cat) + label_ids = nested_xla_mesh_reduce(label_ids, "eval_label_ids") if eval_losses is not None: eval_losses = xm.mesh_reduce("eval_losses", torch.tensor(eval_losses), torch.cat).tolist()