Need a test that checks that the weights can load after training (see #21 )
Need a test that checks that the weights can load after training (see #21 )