extent test on AMD GPU

gher-uliege · Nov 28, 2024 · 621a6dd · 621a6dd
1 parent f05eb6c
commit 621a6dd
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/test/test_reconstruct_points.jl b/test/test_reconstruct_points.jl
@@ -1,10 +1,11 @@
-ENV["CUDA_VISIBLE_DEVICES"]=""
+#ENV["CUDA_VISIBLE_DEVICES"]=""
 
 using DINCAE
 using NCDatasets
 using Random
 using Test
 using CUDA
+using AMDGPU
 
 T = Float32
 filename = "subset-sla-train.nc"
@@ -58,6 +59,8 @@ Random.seed!(seed)
 Atype =
     if CUDA.functional()
         CuArray{T}
+    elseif AMDGPU.functional()
+        ROCArray{T}
     else
         Array{T}
     end

diff --git a/test/test_reconstruct_points_laplacian.jl b/test/test_reconstruct_points_laplacian.jl
@@ -5,6 +5,7 @@ using NCDatasets
 using Random
 using Test
 using CUDA
+using AMDGPU
 
 T = Float32
 filename = "subset-sla-train.nc"
@@ -60,6 +61,8 @@ Random.seed!(seed)
 Atype =
     if CUDA.functional()
         CuArray{T}
+    elseif AMDGPU.functional()
+        ROCArray{T}
     else
         Array{T}
     end