fixes

ext/FluxAMDGPUExt/FluxAMDGPUExt.jl

@@ -6,7 +6,7 @@ import Flux
 import Flux: FluxCPUAdaptor, FluxAMDGPUAdaptor, _amd, adapt_storage, fmap
 import Flux: DenseConvDims, Conv, ConvTranspose, conv, conv_reshape_bias
 import NNlib
-
+using MLDataDevices: MLDataDevices
 using AMDGPU
 using Adapt
 using Random

ext/FluxCUDAExt/FluxCUDAExt.jl

@@ -10,6 +10,7 @@ using ChainRulesCore
 using Random
 using Adapt
 import Adapt: adapt_storage
+using MLDataDevices: MLDataDevices
 
 
 const USE_CUDA = Ref{Union{Nothing, Bool}}(nothing)

ext/FluxCUDAExt/functor.jl

@@ -57,5 +57,5 @@ function _cuda(id::Union{Nothing, Int}, x)
 end
 
 function Flux._get_device(::Val{:CUDA}, id::Int)
-    return MLDataUtils.gpu_device(id+1, force=true)
+    return MLDataDevices.gpu_device(id+1, force=true)
 end

ext/FluxMetalExt/FluxMetalExt.jl

@@ -4,7 +4,7 @@ import Flux
 import Flux: FluxCPUAdaptor, FluxMetalAdaptor, _metal, _isleaf, adapt_storage, fmap
 import NNlib
 using ChainRulesCore
-
+using MLDataDevices: MLDataDevices
 using Metal
 using Adapt
 using Random

ext/FluxMetalExt/functor.jl

@@ -35,6 +35,6 @@ end
 
 function Flux._get_device(::Val{:Metal}, id::Int)
     @assert id == 0 "Metal backend only supports one device at the moment"
-    return MLDataDevices.gpu_device()
+    return MLDataDevices.gpu_device(force=true)
 end
 

test/train.jl

@@ -155,7 +155,7 @@ for (trainfn!, name) in ((Flux.train!, "Zygote"), (train_enzyme!, "Enzyme"))
       pen2(x::AbstractArray) = sum(abs2, x)/2
       opt = Flux.setup(Adam(0.1), model)
 
-      @test_broken begin
+      @test begin
         trainfn!(model, data, opt) do m, x, y
           err = Flux.mse(m(x), y)
           l2 = sum(pen2, Flux.params(m))
@@ -166,7 +166,7 @@ for (trainfn!, name) in ((Flux.train!, "Zygote"), (train_enzyme!, "Enzyme"))
         @test diff1 ≈ diff2
 
         true
-      end
+      end broken = VERSION >= v"1.11"
     end
 
     # Take 3: using WeightDecay instead. Need the /2 above, to match exactly.