From 67a3920d857de520d7354a4d7049ab6f91957340 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9mentine=20Fourrier?=
 <22726840+clefourrier@users.noreply.github.com>
Date: Thu, 2 Feb 2023 16:29:13 +0100
Subject: [PATCH] Fix Graphormer test suite (#21419)

* [FIX] path for Graphormer checkpoint

* [FIX] Test suite for graphormer

* [FIX] Update graphormer default num_classes
---
 .../graphormer/configuration_graphormer.py       |  6 +++---
 .../graphormer/test_modeling_graphormer.py       | 16 ++++++----------
 2 files changed, 9 insertions(+), 13 deletions(-)

diff --git a/src/transformers/models/graphormer/configuration_graphormer.py b/src/transformers/models/graphormer/configuration_graphormer.py
index 17791bfb0c..36a9262b07 100644
--- a/src/transformers/models/graphormer/configuration_graphormer.py
+++ b/src/transformers/models/graphormer/configuration_graphormer.py
@@ -39,8 +39,8 @@ class GraphormerConfig(PretrainedConfig):
 
 
     Args:
-        num_classes (`int`, *optional*, defaults to 2):
-            Number of target classes or labels, set to 1 if the task is a regression task.
+        num_classes (`int`, *optional*, defaults to 1):
+            Number of target classes or labels, set to n for binary classification of n tasks.
         num_atoms (`int`, *optional*, defaults to 512*9):
             Number of node types in the graphs.
         num_edges (`int`, *optional*, defaults to 512*3):
@@ -134,7 +134,7 @@ class GraphormerConfig(PretrainedConfig):
 
     def __init__(
         self,
-        num_classes: int = 2,
+        num_classes: int = 1,
         num_atoms: int = 512 * 9,
         num_edges: int = 512 * 3,
         num_in_degree: int = 512,
diff --git a/tests/models/graphormer/test_modeling_graphormer.py b/tests/models/graphormer/test_modeling_graphormer.py
index 90698d2781..87bf470655 100644
--- a/tests/models/graphormer/test_modeling_graphormer.py
+++ b/tests/models/graphormer/test_modeling_graphormer.py
@@ -40,7 +40,7 @@ class GraphormerModelTester:
     def __init__(
         self,
         parent,
-        num_classes=2,
+        num_classes=1,
         num_atoms=512 * 9,
         num_edges=512 * 3,
         num_in_degree=512,
@@ -614,7 +614,7 @@ class GraphormerModelIntegrationTest(unittest.TestCase):
                     [3, 3, 4, 3, 3, 3, 3, 4, 4, 3, 4, 2, 2, 0, 0, 0, 0],
                 ]
             ),
-            "x": tensor(
+            "input_nodes": tensor(
                 [
                     [[3], [3], [3], [3], [3], [3], [3], [3], [3], [3], [3], [3], [3], [3], [3], [3], [3]],
                     [[3], [3], [3], [3], [3], [3], [3], [3], [3], [3], [3], [3], [3], [0], [0], [0], [0]],
@@ -1279,15 +1279,11 @@ class GraphormerModelIntegrationTest(unittest.TestCase):
 
         output = model(**model_input)["logits"]
 
-        print(output.shape)
-        print(output)
-
-        expected_shape = torch.Size(())
+        expected_shape = torch.Size((2, 1))
         self.assertEqual(output.shape, expected_shape)
 
-        # TODO Replace values below with what was printed above.
-        expected_slice = torch.tensor(
-            [[[-0.0483, 0.1188, -0.0313], [-0.0606, 0.1435, 0.0199], [-0.0235, 0.1519, 0.0175]]]
+        expected_logs = torch.tensor(
+            [[7.6060], [7.4126]]
         )
 
-        self.assertTrue(torch.allclose(output[:, :3, :3], expected_slice, atol=1e-4))
+        self.assertTrue(torch.allclose(output, expected_logs, atol=1e-4))