FromTorchDataset

Contents

FromTorchDataset#

class stable_pretraining.data.FromTorchDataset(dataset, names, transform=None, add_sample_idx=True)[source]#

Bases: Dataset

Wrapper that converts a positional-return PyTorch dataset into dict samples.

Parameters:

  • dataset – PyTorch dataset to wrap.

  • names – List of names for each element returned by the dataset.

  • transform – Optional transform applied to the dict sample.

  • add_sample_idx – If True, adds a sample_idx field to each sample.

Contents