contrete type parameters; remove closures

Author: CarloLucibello

Project.toml

@@ -8,6 +8,7 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
+FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"

src/GraphNeuralNetworks.jl

@@ -12,6 +12,11 @@ using NNlib: scatter, gather
 using ChainRulesCore
 using Reexport
 
+# use `@closure` in conv layers in ored to avoid issues
+# https://github.com/JuliaLang/julia/issues/15276
+# https://github.com/FluxML/Zygote.jl/issues/1317
+# using FastClosures: @closure 
+
 using SparseArrays, Graphs # not needed but if removed Documenter will complain
 
 include("GNNGraphs/GNNGraphs.jl")