batch_to_device¶
- torchtune.utils.batch_to_device(batch: dict, device: device) None[source]¶
Function that takes a dictionary (or nested dictionary) of tensors and sets them all to the same device. This utility is intended to be used for batches of data to be moved to device, the update is inplace.
- Parameters:
batch (dict) – dict of Tensors or more nested dicts of tensors.
device (torch.device) – torch device to move the tensor’s too
- Raises:
AttributeError – if batch dict contains anything other than tensors
