Fixed the test hopefully

FluxML · Feb 11, 2020 · b6455c6 · b6455c6
1 parent bede06d
commit b6455c6
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/test/layers/normalisation.jl b/test/layers/normalisation.jl
@@ -207,7 +207,11 @@ end
       Δg = gswn[g]
       Δv = gswn[v]
       @test wnd(fake_data) ≈ d(fake_data)
-      @test sum(ΔW .* v ./ normv, dims = WN_dim...) ≈ Δg
+      if isa(WN_dim, Int)
+        @test sum(ΔW .* v ./ normv, dims = WN_dim) ≈ Δg
+      else
+        @test sum(ΔW .* v ./ normv, dims = WN_dim[1]) ≈ Δg
+      end
       @test g ./ normv .* ΔW - g .* Δg .* v ./ (normv.^2) ≈ Δv
       @test size(Δv) == size(ΔW)
       @test isa(wnd.layer.W, Flux.WeightNormParam)