diff --git a/custom_utils/custom_utils.py b/custom_utils/custom_utils.py
index 302b5ff..887385c 100644
--- a/custom_utils/custom_utils.py
+++ b/custom_utils/custom_utils.py
@@ -35,10 +35,13 @@ class UV2Mesh():
         # self.y2 = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
         # self.z2 = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32) * 2), axis=-1)
 
-    # input = channels_first
     @tf.function
     def uv_to_mesh(self, uv):
-        uv = tf.transpose(uv, [0,2,3,1])
+        if uv.shape[1] == 3:
+            uv = tf.transpose(uv, [0,2,3,1])
+        if uv.shape[1] != 512:
+            uv = tf.image.resize(uv, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+            
         batch_size = uv.shape[0]
 
         x1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, 14475, 1), dtype=tf.int32)), axis=-1)