Update collating_graphormer.py (#23862)

0623f08e · Clémentine Fourrier · GitHub · 62ba64b9 · 0623f08e
Unverified Commit 0623f08e authored May 30, 2023 by Clémentine Fourrier Committed by GitHub May 30, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

src/transformers/models/graphormer/collating_graphormer.py src/transformers/models/graphormer/collating_graphormer.py +1 -1

No files found.
--- a/src/transformers/models/graphormer/collating_graphormer.py
+++ b/src/transformers/models/graphormer/collating_graphormer.py
@@ -129,6 +129,6 @@ class GraphormerDataCollator:
            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