make the permute faster by doing it on batch

basedformer/dataset.py

@@ -85,6 +85,7 @@ class ShardedImageDataset(data.Dataset):
         #make sure these operations are fast!
         ids = [t[1] for t in tensors]
         tensors = torch.stack([t[0] for t in tensors])
+        tensors = tensors.permute(0, 3, 1, 2).float()
         #####################################
         if self.outer_transform:
             tensors = self.outer_transform(tensors)
@@ -95,7 +96,7 @@ class ShardedImageDataset(data.Dataset):
         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)
+        data = torch.from_numpy(data)#.permute(2, 0, 1)
         if self.inner_transform:
             data = self.inner_transform(data)