Update collating_graphormer.py (#23862)

huggingface · May 30, 2023 · 0623f08 · 0623f08
1 parent 62ba64b
commit 0623f08
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