Model is significantly memory intensive

[yhteoh]

For the small model, I am training with a batch_size of 32 only and sampling 16x16 configurations constantly runs out of memory (for sample size 32) when ran on a graphics card with 6GB VRAM.

Issue #4 may or may not be related.

[therealdavidos]

Inspect model during training, inference and look for bottlenecks.