Regression with ViT

[alemelis]

hello 👋

I'm experimenting with a regression problem in which I'd like to predict 6 floats given an image. I have a working pipeline using a ResNet backbone and a linear layer with leaky activation as output. Loss is computed as RMSE

Now I'd like to see how vision transformers perform on the same task. The approach I'm following is to load a pretrained model as

self.vit = timm.create_model('deit_tiny_patch16_224', pretrained=True, num_classes=6)

wich is used in the forward pass as

y = self.vit(x)

Loss is still computed as RMSE.
I'm using pytorch-lightning to find the optimal learning rate which is found to be 2e-5.
During training the loss doesn't change, hence the training is not converging.

Is this the right way of setting the regression problem?

Many thanks!