fix: Handle multiple heads

Author: lenianiva

GNNlib/src/layers/conv.jl

@@ -139,7 +139,7 @@ function gat_conv(l, g::AbstractGNNGraph, x, e::Union{Nothing, AbstractMatrix} =
     α = dropout(α, l.dropout)
     β = α .* m.Wxj
     x = aggregate_neighbors(g, +, β)
-    width = size(x, 1)
+    width = size(x, 1) * size(x, 2)
 
     if !l.concat
         x = mean(x, dims = 2)
@@ -193,10 +193,11 @@ function gatv2_conv(l, g::AbstractGNNGraph, x, e::Union{Nothing, AbstractMatrix}
     β = α .* m.Wxj
     x = aggregate_neighbors(g, +, β)
 
+    width = size(x, 1) * size(x, 2)
     if !l.concat
         x = mean(x, dims = 2)
     end
-    x = reshape(x, :, size(x, 3))
+    x = reshape(x, width, size(x, 3))
     x = l.σ.(x .+ l.bias)
     return x
 end