InstanceNorm throws a scalar getindex disallowed error on GPU

[andevellicus]

When using InstanceNorm on a 3D CNN, on the GPU with

    CuArrays.allowscalar(false)

I get the following stack trace (abbreviated)

ERROR: LoadError: scalar getindex is disallowed
Stacktrace:
 [1] assertscalar(::String) at /home/###/.julia/packages/GPUArrays/HGtNV/src/host/indexing.jl:41
 [2] getindex at /home/###/.julia/packages/GPUArrays/HGtNV/src/host/indexing.jl:96 [inlined]
 [3] _getindex at ./abstractarray.jl:1003 [inlined]
 [4] getindex at ./abstractarray.jl:980 [inlined]
 [5] _broadcast_getindex at ./broadcast.jl:597 [inlined]
 [6] _getindex at ./broadcast.jl:628 [inlined]
 [7] _broadcast_getindex at ./broadcast.jl:603 [inlined]
 [8] getindex at ./broadcast.jl:564 [inlined]
 [9] copy at ./broadcast.jl:854 [inlined]
 [10] materialize(::Base.Broadcast.Broadcasted{CuArrays.CuArrayStyle{2},Nothing,Zygote.var"#1774#1777"{CuArrays.var"#68#69"{Float32}},Tuple{CuArrays.CuArray{Float32,2,CuArrays.CuArray{Float32,5,Nothing}}}}) at ./broadcast.jl:820
 [11] broadcast_forward(::Function, ::CuArrays.CuArray{Float32,2,CuArrays.CuArray{Float32,5,Nothing}}) at /home/andevellicus/.julia/packages/Zygote/YeCEW/src/lib/broadcast.jl:181
 [12] adjoint at /home/###/.julia/packages/Zygote/YeCEW/src/lib/broadcast.jl:197 [inlined]
 [13] _pullback at /home/###/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:47 [inlined]
 [14] adjoint at /home/###/.julia/packages/Zygote/YeCEW/src/lib/lib.jl:179 [inlined]
 [15] _pullback at /home/###/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:47 [inlined]
 [16] broadcasted at ./broadcast.jl:1232 [inlined]
 [17] InstanceNorm at /home/###/.julia/packages/Flux/Fj3bt/src/layers/normalise.jl:292 [inlined]

BatchNorm works without a problem.

Flux 0.10.4
Julia 1.4.2
CuArrays 2.2.0

[johnnychen94]

bump

julia> l = gpu(InstanceNorm(8));

julia> X = CUDA.randn(Float32, 32, 32, 8, 2);

julia> Flux.mse(X, l(X))
2.5004668f-11

julia> gradient(()->Flux.mse(X, l(X)), Flux.params(l))
ERROR: scalar getindex is disallowed
Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] assertscalar(::String) at /home/jc/local/.julia/packages/GPUArrays/uaFZh/src/host/indexing.jl:41
 [3] getindex at /home/jc/local/.julia/packages/GPUArrays/uaFZh/src/host/indexing.jl:96 [inlined]
 [4] getindex at /home/jc/local/.julia/packages/GPUArrays/uaFZh/src/host/indexing.jl:117 [inlined]
 [5] _broadcast_getindex at ./broadcast.jl:614 [inlined]
 [6] _getindex at ./broadcast.jl:645 [inlined]
 [7] _broadcast_getindex at ./broadcast.jl:620 [inlined]
 [8] getindex at ./broadcast.jl:575 [inlined]
 [9] copy at ./broadcast.jl:876 [inlined]
 [10] materialize(::Base.Broadcast.Broadcasted{CUDA.CuArrayStyle{2},Nothing,Zygote.var"#1159#1162"{CUDA.var"#861#862"{Float32}},Tuple{CuArray{Float32,2}}}) at ./broadcast.jl:837
 [11] broadcast_forward(::Function, ::CuArray{Float32,2}) at /home/jc/local/.julia/packages/Zygote/C94Hq/src/lib/broadcast.jl:181
 [12] adjoint at /home/jc/local/.julia/packages/Zygote/C94Hq/src/lib/broadcast.jl:193 [inlined]
 [13] _pullback at /home/jc/local/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:47 [inlined]
 [14] adjoint at /home/jc/local/.julia/packages/Zygote/C94Hq/src/lib/lib.jl:179 [inlined]
 [15] _pullback at /home/jc/local/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:47 [inlined]
 [16] broadcasted at ./broadcast.jl:1257 [inlined]
 [17] InstanceNorm at /home/jc/Documents/Julia/Flux.jl/src/layers/normalise.jl:313 [inlined]
 [18] _pullback(::Zygote.Context, ::InstanceNorm{typeof(identity),CuArray{Float32,1},CuArray{Float32,1},Float32}, ::CuArray{Float32,4}) at /home/jc/local/.julia/packages/Zygote/C94Hq/src/compiler/interface2.jl:0
 [19] #3 at ./REPL[11]:1 [inlined]
 [20] _pullback(::Zygote.Context, ::var"#3#4") at /home/jc/local/.julia/packages/Zygote/C94Hq/src/compiler/interface2.jl:0
 [21] pullback(::Function, ::Zygote.Params) at /home/jc/local/.julia/packages/Zygote/C94Hq/src/compiler/interface.jl:172
 [22] gradient(::Function, ::Zygote.Params) at /home/jc/local/.julia/packages/Zygote/C94Hq/src/compiler/interface.jl:53
 [23] top-level scope at REPL[11]:1

(v1.5) pkg> st
Status `/mnt/local/home/jc/.julia/environments/v1.5/Project.toml`
  [c7e460c6] ArgParse v1.1.1
  [6e4b80f9] BenchmarkTools v0.5.0
  [052768ef] CUDA v2.3.0
  [5ae59095] Colors v0.12.4
  [863f3e99] Comonicon v0.10.1
  [31a5f54b] Debugger v0.6.6
  [b4f34e82] Distances v0.9.2
  [53c48c17] FixedPointNumbers v0.8.4
  [587475ba] Flux v0.11.2
  [5fb14364] OhMyREPL v0.5.10
  [91a5bcdd] Plots v1.9.1
  [295af30f] Revise v3.1.9
  [5e47fb64] TestImages v1.3.1

[DhairyaLGandhi]

Is that captured in the cuda tests appropriately