add vits architecture

[themattinthehatt]

The current implementation only allows for vit-b/16 architecture, which can be expensive to train. Would be nice to inlcude the smaller vit-s/16 architecture as an alternative.