added onecold broadcast test

2019-02-04 17:57:34 +05:30 · 2019-02-04 17:57:34 +05:30 · 6654ebfc90
parent 30aa814c4d
commit 6654ebfc90
1 changed files with 6 additions and 0 deletions
--- a/test/cuda/cuda.jl
+++ b/test/cuda/cuda.jl
@ -38,6 +38,12 @@ Flux.back!(sum(l))

 end

+@testset "onecold gpu" begin
+  x = rand(Float32, 10, 3) |> gpu;
+  y = Flux.onehotbatch(1:3, 1:10) |> gpu;
+  @test_nowarn Flux.onecold(x) .== Flux.onecold(y)
+end
+
 if CuArrays.libcudnn != nothing
    @info "Testing Flux/CUDNN"
    include("cudnn.jl")