add extra block, lower feature sizes

auto_encoder_static.py

@@ -42,6 +42,7 @@ def init_wandb():
     artifact.add_file('auto_encoder_static.py')
     artifact.add_file('losses.py')
     artifact.add_file('models/base.py')
+    artifact.add_file('models/resnet.py')
     artifact.add_file('models/cond_encoder.py')
     artifact.add_file('dataset.py')
     artifact.add_file('convnext/models/convnext_tf.py')
@@ -54,7 +55,7 @@ uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_
 
 
 
-# @tf.function
+@tf.function
 def train_step(batch):
     with tf.GradientTape() as tape:
         pred = model([batch['uv_static'], batch['uv_mask']], training=True)
@@ -138,10 +139,11 @@ def train():
         if step == 0:
             to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
 
-        if step % 100:    
-            wandb.log(to_log, step=step, commit=True)
-        else:
-            wandb.log(to_log, step=step, commit=False)
+        # if step % 1000 == 0:    
+        #     wandb.log(to_log, step=step, commit=True)
+        # else:
+            # wandb.log(to_log, step=step, commit=False)
+        wandb.log(to_log, step=step)
 
         if step % 10000 == 0:
             manager_enc.save(checkpoint_number=step)
@@ -187,7 +189,7 @@ def test():
 
 
 if __name__ == '__main__':
-    # init_wandb()
+    init_wandb()
     # allow_memory_growth()
     parser = argparse.ArgumentParser(description='')
     parser.add_argument('--name', type=str)