Updated tests in cudnn.jl

test/cuda/cudnn.jl

@@ -1,48 +1,47 @@
 using Flux, CuArrays, Test
+using Zygote
 trainmode(f, x...) = forward(f, x...)[1]
-#
-# @testset "CUDNN BatchNorm" begin
-#     @testset "4D Input" begin
-#         x = Float64.(collect(reshape(1:12, 2, 2, 3, 1)))
-#         m = BatchNorm(3)
-#         cx = gpu(x)
-#         cm = gpu(m)
-#
-#         y = trainmode(m, x)
-#         cy = trainmode(cm, cx)
-#
-#         # @test cy isa TrackedArray{Float32,4,CuArray{Float32,4}}
-#
-#         @test cpu(data(cy)) ≈ data(y)
-#
-#         g = rand(size(y)...)
-#         Flux.back!(y, g)
-#         Flux.back!(cy, gpu(g))
-#
-#         @test m.γ.grad ≈ cpu(cm.γ.grad)
-#         @test m.β.grad ≈ cpu(cm.β.grad)
-#         @test x.grad ≈ cpu(x.grad)
-#     end
-#
-#     @testset "2D Input" begin
-#         x = TrackedArray(Float64.(collect(reshape(1:12, 3, 4))))
-#         m = BatchNorm(3)
-#         cx = gpu(x)
-#         cm = gpu(m)
-#
-#         y = m(x)
-#         cy = cm(cx)
-#
-#         @test cy isa TrackedArray{Float32,2,CuArray{Float32,2}}
-#
-#         @test cpu(data(cy)) ≈ data(y)
-#
-#         g = rand(size(y)...)
-#         Flux.back!(y, g)
-#         Flux.back!(cy, gpu(g))
-#
-#         @test m.γ.grad ≈ cpu(cm.γ.grad)
-#         @test m.β.grad ≈ cpu(cm.β.grad)
-#         @test x.grad ≈ cpu(x.grad)
-#     end
-# end
+
+@testset "CUDNN BatchNorm" begin
+    @testset "4D Input" begin
+        x = Float64.(collect(reshape(1:12, 2, 2, 3, 1)))
+        m = BatchNorm(3)
+        cx = gpu(x)
+        cm = gpu(m)
+
+        y = trainmode(m, x)
+        cy = trainmode(cm, cx)
+
+        @test cpu(data(cy)) ≈ data(y)
+
+        g = rand(size(y)...)
+        # Flux.back!(y, g)
+        # Flux.back!(cy, gpu(g))
+
+        @test m.γ ≈ cpu(cm.γ)
+        @test m.β ≈ cpu(cm.β)
+        @test x ≈ cpu(x)
+    end
+
+    @testset "2D Input" begin
+        x = Float64.(collect(reshape(1:12, 3, 4)))
+        m = BatchNorm(3)
+        cx = gpu(x)
+        cm = gpu(m)
+
+        y = trainmode(m, x)
+        cy = trainmode(cm, cx)
+
+        @test cy isa CuArray{Float32,2}
+
+        @test cpu(data(cy)) ≈ data(y)
+
+        g = rand(size(y)...)
+        #Flux.back!(y, g)
+        #Flux.back!(cy, gpu(g))
+
+        @test m.γ ≈ cpu(cm.γ)
+        @test m.β ≈ cpu(cm.β)
+        @test x ≈ cpu(x)
+    end
+end