Simple ViT error

[matrix0117]

Hello, I'm trying to apply this to cv, using the code from Simple Vit give me this error. There seems to be an error in this code

[NishchalPrasad]

You need to squeeze the tensor in dimensions 1 and 2. It should work then.