diff --git a/data/dataset.py b/data/dataset.py
index 5da5176..fe5f83d 100644
--- a/data/dataset.py
+++ b/data/dataset.py
@@ -24,10 +24,11 @@ from cloth3d.util import quads2tris
 
 
 class Dataset(Dataset):
-    def __init__(self, folder, folder_rgb, frames_file='frames_all.txt', device='cuda', use_cache=True):
+    def __init__(self, folder, folder_rgb, frames_file='frames_all.txt', device='cuda', use_cache=True, random_view=True):
         self.folder = folder
         self.folder_rgb = folder_rgb
         self.device = device
+        self.random_view = random_view
 
         self.transforms = tv.transforms.ToTensor()
         self.frames = read_frames_file(os.path.join(folder, frames_file))
@@ -52,8 +53,9 @@ class Dataset(Dataset):
 
         path = os.path.join(self.folder, frame.id)
         info = self.reader.read_info(frame.id)
+        zrot = info['zrot'] if not self.random_view else (((np.random.random(1) * 2) -1) * np.pi)[0]
         # mesh_posed = readPC2Frame(os.path.join(path, frame.garment + '_registered.pc16'), frame.frame_nr, True)
-        mesh_posed = self.reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, info['trans'][:,frame.frame_nr], info['zrot'])
+        mesh_posed = self.reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, info['trans'][:,frame.frame_nr], zrot)
         # mesh_posed_no_trans = self.reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, absolute=False)
 
         mesh_mask = np.load(os.path.join(path, frame.garment + '_mask.npy'))
@@ -61,7 +63,7 @@ class Dataset(Dataset):
         faces = mask_faces(self.body_faces, mesh_mask)
         # faces = quads2tris(self.reader.read_garment_topology(frame.id, frame.garment))
         # body = readPC2Frame(os.path.join(path, 'Body_posed.pc16'), frame.frame_nr, True)
-        body, body_faces, root_joint, joints, joints_no_trans = self.reader.read_human(frame.id, frame.frame_nr, info['trans'][:,frame.frame_nr], info['zrot'])
+        body, body_faces, root_joint, joints, joints_no_trans = self.reader.read_human(frame.id, frame.frame_nr, info['trans'][:,frame.frame_nr], zrot)
         # body_no_trans, _, _, _, _ = self.reader.read_human(frame.id, frame.frame_nr, absolute=False)
         body_faces = quads2tris(body_faces)
 
@@ -190,19 +192,19 @@ class Dataset(Dataset):
         # v2i_coverage[v2i[:,0], v2i[:, 1]] = 1
         # ex.imshow(v2i_coverage).show()
 
-        rgb_image = tv.io.read_image(os.path.join(self.folder_rgb, frame.id, f'{frame.frame_nr:04d}.png'), mode=tv.io.ImageReadMode.RGB).to(self.device).to(torch.float32)
+        rgb_image = tv.io.read_image(os.path.join(self.folder_rgb, frame.id, f'{frame.frame_nr:04d}.png'), mode=tv.io.ImageReadMode.RGB).to(torch.float32)
 
-        view = torch.tensor(np.concatenate((np.array([info['zrot']], dtype=np.float32), info['camLoc'].astype(np.float32), info['trans'][:,frame.frame_nr].astype(np.float32)), axis=0), device=self.device)
+        view = torch.tensor(np.concatenate((np.array([zrot], dtype=np.float32), info['camLoc'].astype(np.float32), info['trans'][:,frame.frame_nr].astype(np.float32)), axis=0), device=self.device)
 
         # image = self.transforms(image)
-        data = { 'mesh_posed': mesh_posed,
+        data = { 'mesh_posed': torch.tensor(mesh_posed, device=self.device),
                  'depth_image': depth_image, 
                  'mesh_mask': mesh_mask, 
                  'depth_mask': depth_mask, 
                  'normal_image': normal_image,
                  'depth_mask_garment': depth_mask_garment,
                  'depth_mask_body': depth_mask_body,
-                 'rgb_image': rgb_image,
+                 'rgb_image': rgb_image.to(self.device),
                  'view': view,
                 }