Use MLUtils.ones_like

Also go back to indexing instead of `selectdim` to prevent scalar indexing on the GPU

Author: theabhirath

Project.toml

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

src/layers/embeddings.jl

@@ -59,11 +59,8 @@ 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_like(x)
+    tokens = m.token .* MLUtils.ones_like(x, T, (1, 1, size(x, 3)))
     return hcat(tokens, x)
 end
 

src/vit-based/vit.jl

@@ -53,7 +53,7 @@ function vit(imsize::Dims{2} = (256, 256); inchannels = 3, patch_size::Dims{2} =
                      ViPosEmbedding(embedplanes, npatches + 1),
                      Dropout(emb_dropout),
                      transformer_encoder(embedplanes, depth, nheads; mlp_ratio, dropout),
-                     (pool == :class) ? x -> selectdim(x, 2, 1) : seconddimmean),
+                     (pool == :class) ? x -> x[:, 1, :] : seconddimmean),
                Chain(LayerNorm(embedplanes), Dense(embedplanes, nclasses, tanh_fast)))
 end