Testing model code tends to be very difficult unless you design your training loop with lots of abstractions and dependency injection which makes the code less explicit and difficult to understand. For example, try to look at the Tensorflow Estimator framework. Absolutely awful to use but is well tested.