Add Flux.gpu function for TemporalSnapshotsGNNGraph type (#362)

* Add `TemporalGraphConv`

* Improve names

* Add function to move TSG to gpu

* Add `Flux.`

* Add test for `gpu` movement

* Revert "Improve names"

This reverts commit e8532ab.

* Revert "Add `TemporalGraphConv`"

This reverts commit 47b78f1.

* Add test for `gpu` movement

* Add functor macro for `TemporalSnapshotsGNNGraph`

* Add `TEST_GPU`

Author: aurorarossi

src/GNNGraphs/temporalsnapshotsgnngraph.jl

@@ -36,10 +36,10 @@ TemporalSnapshotsGNNGraph:
 ```
 """
 struct TemporalSnapshotsGNNGraph
-    num_nodes::Vector{Int}   
-    num_edges::Vector{Int}
+    num_nodes::AbstractVector{Int}   
+    num_edges::AbstractVector{Int}
     num_snapshots::Int
-    snapshots::Vector{<:GNNGraph}
+    snapshots::AbstractVector{<:GNNGraph}
     tgdata::DataStore   
 end
 
@@ -239,4 +239,6 @@ function print_feature_t(io::IO, feature)
     else 
         print(io, "no")
     end
-end
+end
+
+@functor TemporalSnapshotsGNNGraph

test/GNNGraphs/temporalsnapshotsgnngraph.jl

@@ -101,3 +101,17 @@ end
     tsg.tgdata.x=rand(4)
     @test sprint(show,tsg) == "TemporalSnapshotsGNNGraph(5) with x: 4-element data"
 end
+
+if TEST_GPU
+    @testset "gpu" begin
+        snapshots = [rand_graph(10, 20; ndata = rand(5,10)) for i in 1:5]
+        tsg = TemporalSnapshotsGNNGraph(snapshots)
+        tsg.tgdata.x = rand(5)
+        tsg = Flux.gpu(tsg)
+        @test tsg.snapshots[1].ndata.x isa CuArray
+        @test tsg.snapshots[end].ndata.x isa CuArray
+        @test tsg.tgdata.x isa CuArray
+        @test tsg.num_nodes isa CuArray
+        @test tsg.num_edges isa CuArray
+    end
+end