return ids for real

4577aa6f · novelailab · 8056efa1 · 4577aa6f
Commit 4577aa6f authored Jul 04, 2022 by novelailab
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

basedformer/dataset.py basedformer/dataset.py +1 -1

No files found.
--- a/basedformer/dataset.py
+++ b/basedformer/dataset.py
@@ -90,7 +90,7 @@ class ShardedImageDataset(data.Dataset):
        if self.outer_transform:
            tensors = self.outer_transform(tensors)
            
-        return tensors
+        return tensors, ids

    def read_from_metadata_key(self, key):
        offset, size, id = self.metadata[key]