Update conv.jl: self loops

JuliaGraphs · Sep 30, 2024 · be04df5 · be04df5
1 parent fea71c9
commit be04df5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/GNNLux/src/layers/conv.jl b/GNNLux/src/layers/conv.jl
@@ -929,11 +929,11 @@ 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.heads, l.concat, l.skip_connection, l.add_self_loops)
     return GNNlib.transformer_conv(m, g, x, e), st
 end
 
 function Base.show(io::IO, l::TransformerConv)
     (in, ein), out = (l.in_dims, l.out_dims)
     print(io, "TransformerConv(($in, $ein) => $out, heads=$(l.heads))")
-end
+end