fix buildkite

Author: CarloLucibello

GNNGraphs/test/gnngraph.jl

@@ -98,7 +98,7 @@ end
                 # See https://github.com/JuliaGPU/CUDA.jl/pull/1093
                 mat_gpu = adjacency_matrix(g_gpu)
                 @test mat_gpu isa AbstractMatrix{Int}
-                @test get_device(mat_gpu) == dev
+                @test get_device(mat_gpu) isa AbstractGPUDevice
                 @test Array(mat_gpu) == adj_mat
             end
         end
@@ -108,7 +108,7 @@ end
             if TEST_GPU && !(dev isa MetalDevice)
                 mat_gpu = normalized_laplacian(g_gpu)
                 @test mat_gpu isa AbstractMatrix{Float32}
-                @test get_device(mat_gpu) == dev
+                @test get_device(mat_gpu)isa AbstractGPUDevice
                 @test Array(mat_gpu) == mat
             end
         end
@@ -118,7 +118,7 @@ end
                 mat = scaled_laplacian(g)
                 mat_gpu = scaled_laplacian(g_gpu)
                 @test mat_gpu isa AbstractMatrix{Float32}
-                @test get_device(mat_gpu) == dev
+                @test get_device(mat_gpu) isa AbstractGPUDevice
                 @test Array(mat_gpu) ≈ mat
             end 
         end
@@ -133,10 +133,10 @@ end
                 s_cpu, t_cpu = edge_index(g)
                 s_gpu, t_gpu = edge_index(g_gpu)
                 @test s_gpu isa AbstractVector{Int}
-                @test get_device(s_gpu) == dev
+                @test get_device(s_gpu) isa AbstractGPUDevice
                 @test Array(s_gpu) == s_cpu
                 @test t_gpu isa AbstractVector{Int}
-                @test get_device(t_gpu) == dev
+                @test get_device(t_gpu) isa AbstractGPUDevice
                 @test Array(t_gpu) == t_cpu
             end
         end

GNNGraphs/test/query.jl

@@ -185,7 +185,7 @@ end
             d = degree(g)
             d_gpu = degree(g_gpu)
             @test d_gpu isa AbstractVector{Int}
-            @test get_device(d_gpu) == dev
+            @test get_device(d_gpu) isa AbstractGPUDevice
             @test Array(d_gpu) == d
         end
 
@@ -199,7 +199,7 @@ end
             d = degree(g)
             d_gpu = degree(g_gpu)
             @test d_gpu isa AbstractArray{Float32}
-            @test get_device(d_gpu) == dev
+            @test get_device(d_gpu) isa AbstractGPUDevice
             @test Array(d_gpu) ≈ d
         end
     end

GNNGraphs/test/temporalsnapshotsgnngraph.jl

@@ -145,11 +145,11 @@ end
     tsg.tgdata.x = rand(Float32, 5)
     tsg = tsg |> dev
     @test tsg.snapshots[1].ndata.x isa AbstractMatrix{Float32}
-    @test get_device(tsg.snapshots[1].ndata.x) == dev
+    @test get_device(tsg.snapshots[1].ndata.x) isa AbstractGPUDevice
     @test tsg.snapshots[end].ndata.x isa AbstractMatrix{Float32}
-    @test get_device(tsg.snapshots[end].ndata.x) == dev
+    @test get_device(tsg.snapshots[end].ndata.x) isa AbstractGPUDevice
     @test tsg.tgdata.x isa AbstractVector{Float32}
-    @test get_device(tsg.tgdata.x) == dev
+    @test get_device(tsg.tgdata.x) isa AbstractGPUDevice
     # num_nodes and num_edges are not copied to the device
     @test tsg.num_nodes isa Vector{Int}
     @test tsg.num_edges isa Vector{Int}

GNNGraphs/test/test_module.jl

@@ -28,6 +28,7 @@ using FiniteDifferences: FiniteDifferences
 using Reexport: @reexport
 using MLUtils: MLUtils
 using Zygote: gradient
+using MLDataDevices: AbstractGPUDevice
 @reexport using SparseArrays
 @reexport using MLDataDevices
 @reexport using Random
@@ -36,7 +37,7 @@ using Zygote: gradient
 @reexport using GNNGraphs
 @reexport using Test
 @reexport using Graphs
-export MLUtils, gradient
+export MLUtils, gradient, AbstractGPUDevice
 export ngradient, GRAPH_TYPES