Closes #191 (closed)
Needs testing on a real dataset because training test is failing only because optimizers are different.
