dataset

dataloader.py

@@ -77,7 +77,8 @@ class RT_Dataset(Dataset):
             self.data = self.data[self.data.state == 'holdout']
         elif mode == 'validation':
             self.data = self.data[self.data.state == 'validation']
-
+        if size is not None:
+            self.data = self.data.sample(size)
 
         print('Padding')
         self.data['sequence'] = self.data['sequence'].str.pad(length, side='right', fillchar='_')