Typo

[aksuhton]

There is a typo in vit.jl where the following function is defined

function (m::MultiHeadAttention)(x::AbstractArray{T, 3}, ps, st) where {T}

More precisely the issue is within, where st_ is defined. I believe we want (minimally)

st_ = (qkv_layer=st_qkv, attention_dropout=st_attention, projection=st_projection)

where I've replaced attention -> attention_dropout. Otherwise Zygote.pullback gets confused when I try to run a training loop.

Apologies if it is taboo to not make a proper PR with mwe and error. Thanks again for all your work.

[avik-pal]

You are right. it should be as you suggested. Can you open a PR?

[aksuhton]

Just did over at the following link. This is my first pull request so please feel free to instruct me if I missed something.

#16