D2-equivariant GATv2.

PiperOrigin-RevId: 603596751
google-deepmind · Feb 2, 2024 · fd7eabf · fd7eabf
1 parent 4fe98f9
commit fd7eabf
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/clrs/_src/processors.py b/clrs/_src/processors.py
@@ -339,7 +339,7 @@ def d2_forward(self,
             adj_mat=adj_mat,
             hidden=hidden
         )
-        emb_values.append(cell_embedding)
+        emb_values.append(cell_embedding[0])
       ret_nodes.append(
           jnp.mean(jnp.stack(emb_values, axis=0), axis=0)
       )