Fix ambiguity problems

LuxDL · Feb 24, 2024 · 5effa9c · 5effa9c
1 parent d33b2e3
commit 5effa9c
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 2 deletions.
diff --git a/ext/LuxDeviceUtilsLuxAMDGPUExt.jl b/ext/LuxDeviceUtilsLuxAMDGPUExt.jl
@@ -50,7 +50,9 @@ function adapt_storage(to::LuxAMDGPUAdaptor, x)
         return x_new
     end
 end
+adapt_storage(::LuxAMDGPUAdaptor{Nothing}, rng::AbstractRNG) = rng
 adapt_storage(::LuxAMDGPUAdaptor, rng::AbstractRNG) = rng
+adapt_storage(::LuxAMDGPUAdaptor{Nothing}, rng::Random.TaskLocalRNG) = AMDGPU.rocrand_rng()
 adapt_storage(::LuxAMDGPUAdaptor, rng::Random.TaskLocalRNG) = AMDGPU.rocrand_rng()
 
 adapt_storage(::LuxCPUAdaptor, rng::AMDGPU.rocRAND.RNG) = Random.default_rng()

diff --git a/ext/LuxDeviceUtilsLuxCUDAExt.jl b/ext/LuxDeviceUtilsLuxCUDAExt.jl
@@ -50,7 +50,9 @@ function adapt_storage(to::LuxCUDAAdaptor, x)
         return x_new
     end
 end
+adapt_storage(::LuxCUDAAdaptor{Nothing}, rng::AbstractRNG) = rng
 adapt_storage(::LuxCUDAAdaptor, rng::AbstractRNG) = rng
+adapt_storage(::LuxCUDAAdaptor{Nothing}, rng::Random.TaskLocalRNG) = CUDA.default_rng()
 adapt_storage(::LuxCUDAAdaptor, rng::Random.TaskLocalRNG) = CUDA.default_rng()
 
 adapt_storage(::LuxCPUAdaptor, rng::CUDA.RNG) = Random.default_rng()

diff --git a/test/amdgpu.jl b/test/amdgpu.jl
@@ -82,7 +82,7 @@ if LuxAMDGPU.functional()
         @test typeof(amdgpu_device.device) <: AMDGPU.HIPDevice
         @test AMDGPU.device_id(amdgpu_device.device) == idx
 
-        ps = ps |> amdgpu_device
+        global ps = ps |> amdgpu_device
         @test ps.weight isa ROCArray
         @test ps.bias isa ROCArray
         @test AMDGPU.device_id(AMDGPU.device(ps.weight)) == idx

diff --git a/test/cuda.jl b/test/cuda.jl
@@ -82,7 +82,7 @@ if LuxCUDA.functional()
         @test typeof(cuda_device.device) <: CUDA.CuDevice
         @test cuda_device.device.handle == (idx - 1)
 
-        ps = ps |> cuda_device
+        global ps = ps |> cuda_device
         @test ps.weight isa CuArray
         @test ps.bias isa CuArray
         @test CUDA.device(ps.weight).handle == idx - 1