Update conv.jl: out dims

CarloLucibello · Sep 30, 2024 · c69849f · c69849f
1 parent a6a11bd
commit c69849f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/GNNLux/src/layers/conv.jl b/GNNLux/src/layers/conv.jl
@@ -928,7 +928,7 @@ function (l::TransformerConv)(g, x, e, ps, st)
     BN2 = l.BN2 === nothing ? nothing : 
             StatefulLuxLayer{true}(l.BN2, ps.BN2, _getstate(st, :BN2))
     m = (; W1, W2, W3, W4, W5, W6, FF, BN1, BN2, l.sqrt_out, 
-           l.heads, l.concat, l.skip_connection, l.add_self_loops, l.in_dims)
+           l.heads, l.concat, l.skip_connection, l.add_self_loops, l.in_dims, l.out_dims)
     return GNNlib.transformer_conv(m, g, x, e), st
 end