Refactored some elements of neural net code