Fix OneHotVector/Matrix performance on GPU

[DhairyaLGandhi]

This PR aims to fix performance of OneHotMatrix performance on the GPU. It additionally adds some missing functionality to their behaviour.

[KristofferC]

Does this help with #582?

[DhairyaLGandhi]

For that, we want to ensure onecold returns a CuArray. I'll add a commit to make it so.

[DhairyaLGandhi]

This should help with #582

[MikeInnes]

LGTM, but are there any tests we can add for this? e.g. returns CuArray on the GPU.

[KristofferC]

@test_nowarn doesn't work anymore.

julia> using Test

julia> @test_nowarn @warn "foo"
┌ Warning: foo
└ @ Main REPL[5]:1

[DhairyaLGandhi]

Oops. Perhaps?

@testset "onecold gpu" begin
     x = zeros(Float32, 10, 3) |> gpu;
     y = Flux.onehotbatch(ones(3), 1:10) |> gpu;
     res = Flux.onecold(x) .== Flux.onecold(y)
     @test res isa CuArray
end

[KristofferC]

That seems to indeed better test what we want to check.

[DhairyaLGandhi]

Thanks bunches!