fixed test for optimizer state in trained model history

Author: grantbuster

tests/training/test_train_gan.py

@@ -110,8 +110,15 @@ def test_train(fp_gen, fp_disc, s_enhance, t_enhance, sample_shape, n_epoch=8):
 
         assert np.allclose(model_params['optimizer']['learning_rate'], lr)
         assert np.allclose(model_params['optimizer_disc']['learning_rate'], lr)
-        assert 'learning_rate_gen' in model.history
-        assert 'learning_rate_disc' in model.history
+        assert 'OptmGen/learning_rate' in model.history
+        assert 'OptmDisc/learning_rate' in model.history
+
+        state_cols_b = ['OptmGen/Adam/m/conv3d/bias:0',
+                        'OptmGen/Adam/m/conv2d_transpose/bias:0']
+        state_cols_k = ['OptmGen/Adam/v/conv3d/kernel:0',
+                        'OptmGen/Adam/v/conv2d_transpose/kernel:0']
+        assert any(col in model.history for col in state_cols_b)
+        assert any(col in model.history for col in state_cols_k)
 
         assert 'config_generator' in loaded.meta
         assert 'config_discriminator' in loaded.meta