Apply suggestions from code review

LuxDL · Nov 6, 2024 · bfb5021 · bfb5021
1 parent 4c3714f
commit bfb5021
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 14 deletions.
diff --git a/lib/MLDataDevices/ext/MLDataDevicesChainRulesExt.jl b/lib/MLDataDevices/ext/MLDataDevicesChainRulesExt.jl
@@ -6,14 +6,10 @@ using ChainRules: OneElement
 
 Adapt.adapt_storage(::CPUDevice, x::OneElement) = x
 
-for Dev in (CUDADevice, AMDGPUDevice, MetalDevice, oneAPIDevice, ReactantDevice)
+for dev in (CUDADevice, AMDGPUDevice, MetalDevice, oneAPIDevice, ReactantDevice,
+        CUDADevice{Nothing}, AMDGPUDevice{Nothing})
     # use `@eval` to avoid ambiguity with adapt_storage(::CUDADevice, ::AbstractArray)
-    @eval Adapt.adapt_storage(to::$Dev, x::OneElement) = Adapt.adapt(to, collect(x))
-end
-
-for Dev in (CUDADevice, AMDGPUDevice)
-    # use `@eval` to avoid ambiguity with adapt_storage(::CUDADevice{Nothing}, ::AbstractArray)
-    @eval Adapt.adapt_storage(to::$Dev{Nothing}, x::OneElement) = Adapt.adapt(to, collect(x))
+    @eval Adapt.adapt_storage(to::$(dev), x::OneElement) = Adapt.adapt(to, collect(x))
 end
 
 end
diff --git a/lib/MLDataDevices/ext/MLDataDevicesZygoteExt.jl b/lib/MLDataDevices/ext/MLDataDevicesZygoteExt.jl
@@ -6,14 +6,10 @@ using Zygote: OneElement
 
 Adapt.adapt_storage(::CPUDevice, x::OneElement) = x
 
-for Dev in (CUDADevice, AMDGPUDevice, MetalDevice, oneAPIDevice, ReactantDevice)
+for dev in (CUDADevice, AMDGPUDevice, MetalDevice, oneAPIDevice, ReactantDevice,
+        CUDADevice{Nothing}, AMDGPUDevice{Nothing})
     # use `@eval` to avoid ambiguity with adapt_storage(::CUDADevice, ::AbstractArray)
-    @eval Adapt.adapt_storage(to::$Dev, x::OneElement) = Adapt.adapt(to, collect(x))
-end
-
-for Dev in (CUDADevice, AMDGPUDevice)
-    # use `@eval` to avoid ambiguity with adapt_storage(::CUDADevice{Nothing}, ::AbstractArray)
-    @eval Adapt.adapt_storage(to::$Dev{Nothing}, x::OneElement) = Adapt.adapt(to, collect(x))
+    @eval Adapt.adapt_storage(to::$(dev), x::OneElement) = Adapt.adapt(to, collect(x))
 end
 
 end