Merge pull request #166 from theabhirath/clstoken-fix

Author: darsnack

src/layers/embeddings.jl

@@ -60,8 +60,8 @@ end
 ClassTokens(dim::Integer; init = Flux.zeros32) = ClassTokens(init(dim, 1, 1))
 
 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 .* MLUtils.ones_like(x, T, (1, 1, size(x, 3)))
+    return hcat(tokens, x)
 end
 
 @functor ClassTokens

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