Use host memory to prevent OOM.

JuliaGPU · Dec 10, 2024 · acc31db · acc31db
1 parent 67bbd41
commit acc31db
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/test/base/random.jl b/test/base/random.jl
@@ -200,7 +200,12 @@ end
 
 @testset "counter overflow" begin
     rng = CUDA.RNG()
-    c = CUDA.zeros(Float16, (64, 32, 512, 32, 64))
-    rand!(rng, c)
-    randn!(rng, c)
+    # we may not be able to allocate over 4GB on the GPU, so use CPU memory
+    #c = CUDA.zeros(Float16, (64, 32, 512, 32, 64))
+    c = Array{Float16}(undef, 64, 32, 512, 32, 64)
+    GC.@preserve c begin
+        dc = unsafe_wrap(CuArray, c)
+        rand!(rng, dc)
+        randn!(rng, dc)
+    end
 end