[Feature] ST Transformer Layer

[alexcbb]

Feature details

One of the most important part of the model is the spatio-temporal transformer layer used in the model. It is a memory efficient version from ViT

This layer combines the idea from the original idea with Vision Transformer.

What needs to be done

• [x] Code the vision transformer layer
• [ ] #5
• [ ] Combine both blocks
• [ ] Evaluate the model on toy scenarios
• [ ] Prepare unit tests
• [ ] PR the new feature