fix GATConv layer

Author: yuehhua

src/GeometricFlux.jl

@@ -78,8 +78,6 @@ include("layers/misc.jl")
 include("sampling.jl")
 include("embedding/node2vec.jl")
 
-include("cuda/conv.jl")
-
 using .Datasets
 
 

src/cuda/conv.jl

@@ -253,8 +253,9 @@ function apply_batch_message(gat::GATConv, i, js, X::AbstractMatrix)
 end
 
 function update_batch_edge(gat::GATConv, sg::SparseGraph, E::AbstractMatrix, X::AbstractMatrix, u)
-    @assert check_self_loops(sg) "a vertex must have self loop (receive a message from itself)."
-    mapreduce(i -> apply_batch_message(gat, i, neighbors(sg, i), X), hcat, 1:nv(sg))
+    @assert Zygote.ignore(() -> check_self_loops(sg)) "a vertex must have self loop (receive a message from itself)."
+    ys = map(i -> apply_batch_message(gat, i, GraphSignals.cpu_neighbors(sg, i), X), 1:nv(sg))
+    return hcat(ys...)
 end
 
 function check_self_loops(sg::SparseGraph)

src/layers/conv.jl

@@ -71,6 +71,13 @@
     end
 
     @testset "GATConv" begin
+        adj = T[1 1 0 1;
+                1 1 1 0;
+                0 1 1 1;
+                1 0 1 1]
+        
+        fg = FeaturedGraph(adj)
+
         gat = GATConv(fg, in_channel=>out_channel) |> gpu
         @test size(gat.weight) == (out_channel, in_channel)
         @test size(gat.bias) == (out_channel,)