levit_c.LeViT_c_384 not working

[aashusri]

Hi, there seems to be a bug in the levit_c.LeViT_c_384 model. Could you please have a look at it?

To reproduce.
import torch
import levit_c
model = levit_c.LeViT_c_384()
out = model(torch.rand(4,3,224,224)) # gives following error
RuntimeError: The size of tensor a (384) must match the size of tensor b (4) at non-singleton dimension 1

[btgraham]

Thanks for the report. It should be working now.