The assignment should be succeeding fine, but it appears scalar indexing is happening in the repl when trying to show the tuple of ytrain, xtest in the first example:
julia> Flux.onehotbatch(1:2, 1:2) |> cpu, Flux.onehotbatch(1:2, 1:2) |> cpu
(Bool[1 0; 0 1], Bool[1 0; 0 1])
julia> Flux.onehotbatch(1:2, 1:2) |> gpu, Flux.onehotbatch(1:2, 1:2) |> gpu
(Bool[Error showing value of type Tuple{Flux.OneHotArray{UInt32, 2, 1, 2, CuArray{UInt32, 1, CUDA.Mem.DeviceBuffer}}, Flux.OneHotArray{UInt32, 2, 1, 2, CuArray{UInt32, 1, CUDA.Mem.DeviceBuffer}}}:
ERROR: Scalar indexing is disallowed.
Invocation of getindex resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore are only permitted from the REPL for prototyping purposes.
If you did intend to index this array, annotate the caller with @allowscalar.
...
I’ve opened https://github.com/FluxML/Flux.jl/issues/1905 to track this.