quanda.utils.datasets.on_device_dataset module

Module to move a dataset to a device.

class quanda.utils.datasets.on_device_dataset.OnDeviceDataset(dataset: Dataset, device: str | device)

Bases: Dataset

Wrapper that moves a dataset’s tensors to a target device.

Handles arbitrary sample structures returned by dataset[i] — single tensor, tuple/list of any length, dict — by recursively moving every tensor leaf while leaving non-tensor values untouched.

__init__(dataset: Dataset, device: str | device)

Construct the OnDeviceDataset class.

Parameters:

• dataset (torch.utils.data.Dataset) – The dataset to move to the device.

• device (Union[str, torch.device]) – The device to move the dataset to.