diff --git a/dataset.py b/dataset.py
index 38f665a..d5dbf8b 100644
--- a/dataset.py
+++ b/dataset.py
@@ -59,7 +59,7 @@ class StaticDataset():
         self.signature = signature
         self.folder = folder
         assert set(resolutions).issubset(set([32, 128, 512]))
-        self.possible_keys = ['uv_static', 'uv_static_32', 'uv_static_128' 'mesh_static', 'mesh_mask', 'uv_mask', 'uv_shaped_body']
+        self.possible_keys = ['uv_static', 'uv_static_32', 'uv_static_128' 'mesh_static', 'mesh_mask', 'uv_mask', 'uv_mask_32', 'uv_mask_128', 'uv_shaped_body']
         self.keys = [value for value in signature.keys() if value in self.possible_keys]
         if 'uv_static_128' in signature.keys():
             self.keys.append('uv_static_128')
@@ -79,14 +79,14 @@ class StaticDataset():
 
             if 'uv_static' + res_string in self.signature.keys():
                 uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id +  '.npz')); uv_static = uv_static['uv_static' + self.filled]; uv_static = np.transpose(uv_static, [2, 0, 1]); data['uv_static' + res_string] = uv_static
+            if 'uv_mask' in self.signature.keys():
+                uv_mask = np.load(os.path.join(self.folder, 'uv_mask' + res_string, id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask' + res_string] = uv_mask
 
 
         if 'mesh_static' in self.signature.keys():
             mesh_static = np.load(os.path.join(self.folder, 'mesh_static', id +  '.npz')); mesh_static = mesh_static['mesh_static']; data['mesh_static'] = mesh_static
         if 'mesh_mask' in self.signature.keys():
             mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32); data['mesh_mask'] = mesh_mask
-        if 'uv_mask' in self.signature.keys():
-            uv_mask = np.load(os.path.join(self.folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask'] = uv_mask
         if 'uv_shaped_body' in self.signature.keys():
             uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; uv_shaped_body = np.transpose(uv_shaped_body, [2, 0, 1]); data['uv_shaped_body'] = uv_shaped_body
 
@@ -143,7 +143,7 @@ class DatasetFactory():
 
     
     def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=True, resolutions=[512]):
-        assert (use_filled and len(resolutions) >= 1) or (not use_filled and len(resolutions) == 1)
+        # assert (use_filled and len(resolutions) >= 1) or (not use_filled and len(resolutions) == 1)
         assert set(resolutions).issubset(set([32, 128, 512]))
 
         # static_data = StaticDataset(folder, signature)