Update train.py

Author: zsdonghao

train.py

@@ -97,8 +97,8 @@ def _map_fn_train(img):
         with tf.GradientTape() as tape:
             fake_hr_patchs = G(lr_patchs)
             mse_loss = tl.cost.mean_squared_error(fake_hr_patchs, hr_patchs, is_mean=True)
-        grad = tape.gradient(mse_loss, G.weights)
-        g_optimizer_init.apply_gradients(zip(grad, G.weights))
+        grad = tape.gradient(mse_loss, G.trainable_weights)
+        g_optimizer_init.apply_gradients(zip(grad, G.trainable_weights))
         step += 1
         epoch = step//n_step_epoch
         print("Epoch: [{}/{}] step: [{}/{}] time: {}s, mse: {} ".format(