Use @non_differentiable function for fill! in ClassTokens

Should solve at least part of #165

Author: theabhirath

Project.toml

@@ -5,6 +5,7 @@ version = "0.7.1"
 [deps]
 Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
 BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 LazyArtifacts = "4af54fe1-eca0-43a8-85a7-787d91b784e3"

src/layers/Layers.jl

@@ -5,6 +5,7 @@ using Flux: outputsize, Zygote
 using Functors
 using Statistics
 using MLUtils
+using ChainRulesCore
 
 include("../utilities.jl")
 

src/layers/embeddings.jl

@@ -59,9 +59,12 @@ end
 
 ClassTokens(dim::Integer; init = Flux.zeros32) = ClassTokens(init(dim, 1, 1))
 
+_fill_like(y::AbstractArray{T, 3}) where {T} = fill!(similar(y, 1, 1, size(y, 3)), one(T))
+ChainRulesCore.@non_differentiable _fill_like(y)
+
 function (m::ClassTokens)(x::AbstractArray{T, 3}) where {T}
-  tokens = m.token .* fill(one(T), (1, 1, size(x, 3)))
-  return hcat(tokens, x)
+    tokens = m.token .* _fill_like(x)
+    return hcat(tokens, x)
 end
 
 @functor ClassTokens