Merge pull request #1959 from FluxML/cl/oh

onehotbatch with CuArray

Author: CarloLucibello

src/onehot.jl

@@ -185,6 +185,8 @@ julia> reshape(1:15, 3, 5) * oh  # this matrix multiplication is done efficientl
 """
 onehotbatch(data, labels, default...) = _onehotbatch(data, length(labels) < 32 ? Tuple(labels) : labels, default...)
 
+_onehotbatch(data::CuArray, labels) = _onehotbatch(data |> cpu, labels) |> gpu
+
 function _onehotbatch(data, labels)
   indices = UInt32[something(_findval(i, labels), 0) for i in data]
   if 0 in indices

test/cuda/cuda.jl

@@ -45,6 +45,13 @@ end
 
   gA = rand(3, 2) |> gpu;
   @test gradient(A -> sum(A * y), gA)[1] isa CuArray
+
+  # construct from CuArray
+  x = [1, 3, 2]
+  y = Flux.onehotbatch(x, 0:3)
+  y2 = Flux.onehotbatch(x |> gpu, 0:3)
+  @test y2.indices isa CuArray
+  @test y2 |> cpu == y
 end
 
 @testset "onecold gpu" begin