diff --git a/auto_encoder_static.py b/auto_encoder_static.py
index 7c3d766..be5857f 100644
--- a/auto_encoder_static.py
+++ b/auto_encoder_static.py
@@ -2,12 +2,24 @@ import os
 from configs.config_ae_static import config
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 
+import tensorflow as tf
+gpus = tf.config.experimental.list_physical_devices('GPU')
+if gpus:
+    try:
+        # Currently, memory growth needs to be the same across GPUs
+        for gpu in gpus:
+            tf.config.experimental.set_memory_growth(gpu, True)
+        logical_gpus = tf.config.experimental.list_logical_devices('GPU')
+        print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
+    except RuntimeError as e:
+        # Memory growth must be set before GPUs have been initialized
+        print(e)
+
 import time
 import argparse
 import wandb
 import random
 import numpy as np
-import tensorflow as tf
 from custom_utils.custom_utils import UV2Mesh, load_model, maskFacesBatch
 import plotly.graph_objects as go
 import plotly.express as ex
@@ -25,7 +37,7 @@ from models.cond_encoder import AutoEncoder
 
 
 def init_wandb():
-    run = wandb.init(project='vae', name=config['name'],
+    run = wandb.init(project='ae', name=config['name'],
             notes='convnext + convnext, add dense with tanh at the end of reverse convnext,\
     flattened latent dim to 128\
     train on 100 garments')
@@ -67,24 +79,19 @@ def train_step(batch):
 
 @tf.function
 def valid_step(batch):
-    pred = model(batch['uv_static'])
-    uv_mask = batch['uv_mask']
+    pred = model([batch['uv_static'], batch['uv_mask']], training=False)
 
     mesh_pred = uv2mesh(pred)
     mesh_orig = batch['mesh_static']
+
     e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
-    c_loss = 0
-    l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_static'] - pred)), axis=1) * uv_mask)
-    ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_static'], [0,2,3,1]) + 1, max_val=2))
-    fft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_static'], tf.zeros_like(batch['uv_static'])))), axis=[1,2,3]))
+    l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_static'] - pred)), axis=1))
+    ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 1, batch['uv_static'] + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
     loss = 1e3 * (1 - ms_ssim)
-    mesh_orig_2 = uv2mesh(batch['uv_static'])
+
     metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
-    metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
-  
-    return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim,
-             'valid/fft_loss': fft_loss, 'valid/euclid_error_reconst': metric_2
-            }
+    return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim, }
+       
 
 
 def plot_everything(batch: dict, type: str):
@@ -100,33 +107,35 @@ def plot_everything(batch: dict, type: str):
 
     F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy")
     mask = batch['mesh_mask'].numpy().astype(bool)
-    garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
+    garment_static, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
     gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
     
     body = gt
-    garment_posed = mesh_pred
+    garment_posed = garment_static
     collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
-    collision_faces = [np.concatenate((F, faces[i] + np.sum(mask[i])), axis=0) for i in range(len(garment_posed))]
+    collision_faces = [np.concatenate((faces[i], faces[i] + np.sum(mask[i])), axis=0) for i in range(len(garment_posed))]
     collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
 
-    figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
+    figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_static[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
     
     return { **figs, **{ type + '/images': imgs } }
 
 
 def train():
+    valid_iterator = iter(dataset_valid)
+
     step = 0
     for batch in dataset_train:
         batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
         train_loss = train_step(batch)
 
-        # data_valid = valid_iterator.get_next()
-        # data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
-        valid_loss = {}#valid_step(data_valid)
+        data_valid = valid_iterator.get_next()
+        data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
+        valid_loss = valid_step(data_valid)
 
-        if step % 2500 == 0:
+        if step % 5000 == 0 or (step < 5000 and step % 500 == 0):
             train_figs = plot_everything(batch, 'train')
-            valid_figs = {}#plot_everything(data_valid, 'valid')
+            valid_figs = plot_everything(data_valid, 'valid')
             wandb.log({ **train_loss, **valid_loss, **train_figs, **valid_figs }, step=step)
         else:
             wandb.log({ **train_loss, **valid_loss }, step=step)
@@ -171,10 +180,10 @@ if __name__ == '__main__':
 
     factory = DatasetFactory()
 
-    dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']])
-    static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']])
-    dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'])
-    static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'])
+    dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'])
+    static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True)
+    dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'])
+    static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True)
 
     model : AutoEncoder = AutoEncoder(config['cond_ae'])
     optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
diff --git a/configs/config_ae_static.py b/configs/config_ae_static.py
index 13b0ed9..0c75b05 100644
--- a/configs/config_ae_static.py
+++ b/configs/config_ae_static.py
@@ -1,8 +1,8 @@
 from configs.config_base import config
 
-config['gpu'] = '2'
+config['gpu'] = '5'
 config['lr'] = 0.0005
-config['name'] = 'ae_static_7_2_skirt'
+config['name'] = 'static_1'
 
 config['cond_ae']['upscale_output'] = False