Refactor train_network.py (and related stuff)

The goal here is to refactor this script and all related implementations, to make them compliant to the current framework