Merge pull request #162 from theabhirath/vit-fix

Author: darsnack

src/layers/embeddings.jl

@@ -60,7 +60,7 @@ 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!(similar(x, 1, 1, size(x, 3)), one(T))
+  tokens = m.token .* fill(one(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 -> x[:, 1, :] : seconddimmean),
+                     (pool == :class) ? x -> selectdim(x, 2, 1) : seconddimmean),
                Chain(LayerNorm(embedplanes), Dense(embedplanes, nclasses, tanh_fast)))
 end