Add more unit tests

[tmabraham]

Here are some tests of interest:

• [ ] Test the loss (example fake and real images for the reconstruction loss and discriminator)

• [ ] Check the batch independence and the model parameter updates

• [ ] Test successful overfitting on a single batch

• [ ] Test for rotation invariance and other invariance properties