Supports models from the example training code. Weight loading will probably have to be changed for native formats.