Update dataset.py

This commit ensures that the log displays the correct dataset size when fine-tuning on different HDF5VLADataset.
logs refer to https://github.com/huangxu1991/RoboticsDiffusionTransformer/blob/main/train/train.py#L335C1-L336C1

Author: huangxu1991

train/dataset.py

@@ -188,7 +188,10 @@ def _load_data_from_chunk(chunk_dir, chunk_item_idx):
         raise RuntimeError("Failed to load sample.")
 
     def __len__(self) -> int:
-        return self.num_chunks * self.chunk_size
+        if self.use_hdf5:
+            return len(self.hdf5_dataset)
+        else:
+            return self.num_chunks * self.chunk_size
 
     def _safe_load(self, index):
         read_chunk_item_indices = []