Add tests

[DhairyaLGandhi]

We need tests here. One of the issues is how to actually test this in a representative scenario since we may not have access to multiple devices and enough data / memory available in CI. I can start with setting up some basic routines to make sure we are able to read files, instantiate models, run forward/backward passes etc correctly but those are less interesting since those are better tested elsewhere.

[DhairyaLGandhi]

We can close this once we add some tests for the data loading.