Update collating_graphormer.py (huggingface#23862)

novice03 · Jun 23, 2023 · c8ff490 · c8ff490
1 parent 22c10aa
commit c8ff490
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/transformers/models/graphormer/collating_graphormer.py b/src/transformers/models/graphormer/collating_graphormer.py
@@ -129,6 +129,6 @@ def __call__(self, features: List[dict]) -> Dict[str, Any]:
             else:  # binary classification
                 batch["labels"] = torch.from_numpy(np.concatenate([i["labels"] for i in features]))
         else:  # multi task classification, left to float to keep the NaNs
-            batch["labels"] = torch.from_numpy(np.stack([i["labels"] for i in features], dim=0))
+            batch["labels"] = torch.from_numpy(np.stack([i["labels"] for i in features], axis=0))
 
         return batch