Test for potential multi-device

LuxDL · Jun 7, 2024 · 82a3315 · 82a3315
1 parent e00224b
commit 82a3315
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 0 deletions.
diff --git a/test/amdgpu.jl b/test/amdgpu.jl
@@ -93,6 +93,17 @@ using FillArrays, Zygote  # Extensions
 
     ps_mixed = (; a=rand(2), b=device(rand(2)))
     @test_throws ArgumentError get_device(ps_mixed)
+
+    dev = gpu_device()
+    x = rand(Float32, 10, 2)
+    x_dev = x |> dev
+    @test get_device(x_dev) isa parameterless_type(typeof(dev))
+
+    if LuxDeviceUtils.functional(LuxAMDGPUDevice)
+        dev2 = gpu_device(length(AMDGPU.devices()))
+        x_dev2 = x_dev |> dev2
+        @test get_device(x_dev2) isa typeof(dev2)
+    end
 end
 
 @testset "Wrapped Arrays" begin

diff --git a/test/cuda.jl b/test/cuda.jl
@@ -92,6 +92,17 @@ using FillArrays, Zygote  # Extensions
 
     ps_mixed = (; a=rand(2), b=device(rand(2)))
     @test_throws ArgumentError get_device(ps_mixed)
+
+    dev = gpu_device()
+    x = rand(Float32, 10, 2)
+    x_dev = x |> dev
+    @test get_device(x_dev) isa parameterless_type(typeof(dev))
+
+    if LuxDeviceUtils.functional(LuxCUDADevice)
+        dev2 = gpu_device(length(CUDA.devices()))
+        x_dev2 = x_dev |> dev2
+        @test get_device(x_dev2) isa typeof(dev2)
+    end
 end
 
 @testset "Wrapped Arrays" begin