prototype pipeline

data/dataset.py

@@ -1,6 +1,8 @@
 import os
 import sys
 
+os.environ["CUDA_VISIBLE_DEVICES"] = '1'
+
 import math
 import torch
 import torchvision as tv
@@ -18,14 +20,16 @@ class Dataset(Dataset):
         return len(self.frames)
 
     def __getitem__(self, idx):
-        img_path = os.path.join(self.folder, 'mesh_posed', self.frames[idx])
-        image = np.load(img_path)['mesh_posed']
+        path = os.path.join(self.folder, 'mesh_posed', self.frames[idx])
+        mesh_posed = np.load(path)['mesh_posed']
         # image = self.transforms(image)
-        return image
+        return { 'mesh_posed': mesh_posed }
 
 
 if __name__ == '__main__':
     from torch.utils.data import DataLoader
+    
+    device = "cuda" if torch.cuda.is_available() else "cpu"
 
     dataset = Dataset('/home/workspace/local-data/train')
     train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, num_workers=2)