return identity

basedformer/dataset.py

@@ -82,21 +82,24 @@ class ShardedImageDataset(data.Dataset):
         keys = [*range(key, key+self.bsz)]
         tensors = self.executor.map(self.read_from_metadata_key, keys)
         tensors = list(tensors)
-        tensors = torch.stack(tensors)
+        #make sure these operations are fast!
+        ids = [t[1] for t in tensors]
+        tensors = torch.stack([t[0] for t in tensors])
+        #####################################
         if self.outer_transform:
             tensors = self.outer_transform(tensors)
             
         return tensors
 
     def read_from_metadata_key(self, key):
-        offset, size, d_id = self.metadata[key]
+        offset, size, id = self.metadata[key]
         data = self.mmap[offset:offset+size]
         data = decode_jpeg(data)
         data = torch.from_numpy(data).permute(2, 0, 1)
         if self.inner_transform:
             data = self.inner_transform(data)
 
-        return data
+        return data, id
 
 class ImageDatasetBuilder():
     def __init__(self, folder_path, name, threads=None):