lot more training

cvae_pyramid.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from configs.iccv.config_abl_pyr_32_64_128 import config
+from configs.iccv.config_abl_pyr_32_64_reg_skirt import config
 
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 import tensorflow as tf
@@ -212,7 +212,7 @@ def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, col
     return fig
 
 
-# @tf.function
+@tf.function
 def train_step(batch):
     levels_input, conds = group_inputs_for_levels(batch, config)