Scale our Network and Improve Training

[AnimeshSinha1309]

We need accuracy to increase, by means not related to our pooling module.

• [ ] Add a learning rate scheduler, possibly cyclic / cosine
• [x] Implement a temperature scheduler for Gumbel softmax
• [x] Add residual connections to stop over-smoothing
• [x] Add batch normalization (or other) in the network
• [ ] Figure out how to sample and stabilize, or why it is not working (with math insights)