trying to get resnet working

auto_encoder_static.py

@@ -37,7 +37,7 @@ from models.cond_encoder import AutoEncoderNew as AutoEncoder
 
 
 def init_wandb():
-    run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp', tags=config['tags'])
+    run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/tmp', tags=config['tags'])
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('auto_encoder_static.py')
     artifact.add_file('losses.py')
@@ -55,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)
@@ -139,10 +139,7 @@ def train():
         if step == 0:
             to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
 
-        # 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:
@@ -188,8 +185,20 @@ def test():
     print()
 
 
+class MyLRSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
+
+  def __init__(self, initial_learning_rate):
+    self.initial_learning_rate = initial_learning_rate
+
+  def __call__(self, step):
+    if step < 1000:
+        return self.initial_learning_rate / 10.0
+    else:
+        return self.initial_learning_rate
+  
+
 if __name__ == '__main__':
-    # init_wandb()
+    init_wandb()
     parser = argparse.ArgumentParser(description='')
     parser.add_argument('--name', type=str)
     args = vars(parser.parse_args())