Allow non-square images for terratorch VIT models.

[zaheedgaff]

Describe the issue Currently, the terratorch framework for VIT models, does not allow non-square images to be input into the training of models. However, as discussed with the developers, this restriction can be relaxed.

To Reproduce (optional, but appreciated) Steps to reproduce the behavior:

1. Generate images of non-square dimensions (for example: 208x258)
2. Configure a PyTorch lightning task using the config.yaml file. Use the above data, and use the pre-trained Prithvi 100 model as the backbone.
3. Submit fine-tuning job to cluster/machine
4. See error

Screenshots or log output (optional) If applicable, add screenshots or log output to help explain your problem.

Log Output

Error while processing rearrange-reduction pattern "batch (t h w) e -> batch (t e) h w".
 Input tensor shape: torch. Size([8, 208, 768]). Additional info: {'e': 768, 't': 1, 'h': 14}.
 Shape mismatch, can't divide axis of length 208 in chunks of 14

Expected behavior (optional) The expected behavior is that the fine-tuning process takes place, learning from non-square images.

Deployment information (optional) Describe what you've deployed and how:

• TerraTorch version: 0.99.3
• Installation source: pip

[singam96]

Please check this PR #218

[Joao-L-S-Almeida]

Has #218 solved the issue @zaheedgaff ?