[examples] fix Graphormer as key in state_dict has changed

dmlc · Dec 21, 2023 · 08bf92d · 08bf92d
1 parent 77bf0ac
commit 08bf92d
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/examples/core/Graphormer/README.md b/examples/core/Graphormer/README.md
@@ -24,7 +24,7 @@ How to run
 ----------
 
 ```bash
-accelerate launch --multi_gpu --mixed_precision=fp16 train.py
+accelerate launch --multi_gpu --mixed_precision=fp16 main.py
 ```
 > **_NOTE:_**  The script will automatically download weights pre-trained on PCQM4Mv2. To reproduce the same result, set the total batch size to 64.
 

diff --git a/examples/core/Graphormer/model.py b/examples/core/Graphormer/model.py
@@ -47,7 +47,7 @@ def __init__(
         self.spatial_encoder = SpatialEncoder(
             max_dist=num_spatial, num_heads=num_attention_heads
         )
-        self.graph_token_virtual_dist = nn.Embedding(1, num_attention_heads)
+        self.graph_token_virtual_distance = nn.Embedding(1, num_attention_heads)
 
         self.emb_layer_norm = nn.LayerNorm(self.embedding_dim)
 
@@ -112,7 +112,9 @@ def forward(
         attn_bias[:, 1:, 1:, :] = path_encoding + spatial_encoding
 
         # spatial encoding of the virtual node
-        t = self.graph_token_virtual_dist.weight.reshape(1, 1, self.num_heads)
+        t = self.graph_token_virtual_distance.weight.reshape(
+            1, 1, self.num_heads
+        )
         # Since the virtual node comes first, the spatial encodings between it
         # and other nodes will fill the 1st row and 1st column (omit num_graphs
         # and num_heads dimensions) of attn_bias matrix by broadcasting.