diff --git a/dataset.py b/dataset.py
index 2fff2a5..7277d28 100644
--- a/dataset.py
+++ b/dataset.py
@@ -17,7 +17,7 @@ class DatasetSimple(tf.data.Dataset):
         for frame in frames:
             uv = np.load(os.path.join(uv_folder, frame))
             uv = uv['uv_gt']
-            uv = cv2.resize(uv, [64, 128])
+            # uv = cv2.resize(uv, [64, 128])
             uv = np.transpose(uv, [2, 0, 1])
 
             pose_shape = np.load(os.path.join(pose_hape_folder, frame.split('.')[0] + '.npy'))
diff --git a/train.py b/train.py
index 8f02526..20d2a96 100644
--- a/train.py
+++ b/train.py
@@ -1,6 +1,7 @@
 import os
 import time
 import argparse
+import wandb
 from cv2 import drawFrameAxes
 import numpy as np
 import tensorflow as tf
@@ -186,14 +187,14 @@ class Trainer(object):
         return g_loss, g_gan_loss, pl_penalty
 
     def train(self, dist_datasets, strategy):
-        # @tf.function
+        @tf.function
         def dist_d_train_step(inputs):
             per_replica_losses = strategy.run(fn=self.d_train_step, args=(inputs,))
             mean_d_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
             mean_cond_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
             return mean_d_loss, mean_cond_loss
 
-        # @tf.function
+        @tf.function
         def dist_d_train_step_reg(inputs):
             per_replica_losses = strategy.run(fn=self.d_train_step_reg, args=(inputs,))
             mean_d_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
@@ -202,13 +203,13 @@ class Trainer(object):
             mean_cond_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None)
             return mean_d_loss, mean_d_gan_loss, mean_r1_penalty, mean_cond_loss
 
-        # @tf.function
+        @tf.function
         def dist_g_train_step(inputs):
             per_replica_losses = strategy.run(fn=self.g_train_step, args=(inputs,))
             mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)
             return mean_g_loss
 
-        # @tf.function
+        @tf.function
         def dist_g_train_step_reg(inputs):
             per_replica_losses = strategy.run(fn=self.g_train_step_reg, args=(inputs,))
             mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
@@ -361,9 +362,9 @@ def main():
     parser.add_argument('--use_custom_cuda', type=str_to_bool, nargs='?', const=True, default=True)
     parser.add_argument('--model_base_dir', default='./models', type=str)
     parser.add_argument('--tfrecord_dir', default='./tfrecords', type=str)
-    parser.add_argument('--train_res', default=128, type=int)
+    parser.add_argument('--train_res', default=512, type=int)
     parser.add_argument('--shuffle_buffer_size', default=1000, type=int)
-    parser.add_argument('--batch_size_per_replica', default=1, type=int)
+    parser.add_argument('--batch_size_per_replica', default=12, type=int)
     args = vars(parser.parse_args())
 
     # check tensorflow version
@@ -377,12 +378,12 @@ def main():
 
     # network params
     resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
-    # featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
-    featuremaps = [64, 64, 64, 64, 64, 32, 16, 8, 4]
+    featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
+    # featuremaps = [64, 64, 64, 64, 64, 32, 16, 8, 4]
     train_resolutions, train_featuremaps = filter_resolutions_featuremaps(resolutions, featuremaps, args['train_res'])
     g_params = {
-        'z_dim': 16,
-        'w_dim': 16,
+        'z_dim': 512,
+        'w_dim': 512,
         'labels_dim': 83,
         'n_mapping': 8,
         'resolutions': train_resolutions,
@@ -400,7 +401,7 @@ def main():
 
     # prepare dataset
     # dataset = get_ffhq_dataset(args['tfrecord_dir'], args['train_res'], batch_size=global_batch_size, epochs=None)
-    dataset = DatasetSimple('/workspace/Dataset/stylegan').batch(global_batch_size)#.shuffle(buffer_size=1000).prefetch(buffer_size=tf.data.experimental.AUTOTUNE).cache()
+    dataset = DatasetSimple('/home/worksapce/data/dataset/stylegan').batch(global_batch_size)#.shuffle(buffer_size=1000).prefetch(buffer_size=tf.data.experimental.AUTOTUNE).cache()
     options = tf.data.Options()
     options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
     dataset = dataset.with_options(options)