trw.transforms.transform_move_to_device

Module Contents

Classes

TransformMoveToDevice

Move a tensor to a specified device.

Functions

move_to_device(feature_names: List[str], batch: trw.basic_typing.Batch, device: torch.device, non_blocking: bool) → trw.basic_typing.Batch

trw.transforms.transform_move_to_device.move_to_device(feature_names: List[str], batch: trw.basic_typing.Batch, device: torch.device, non_blocking: bool) → trw.basic_typing.Batch

class trw.transforms.transform_move_to_device.TransformMoveToDevice(device: torch.device, non_blocking: bool = False, criteria_fn: Optional[trw.transforms.transforms.CriteriaFn] = None)

Bases: trw.transforms.transforms.TransformBatchWithCriteria

Move a tensor to a specified device.

Transfert from CPU to GPU can can’t significant time. This transfer time can be masked by transferring the data as part of the data preprocessing on a single GPU system.

Note

This requires to start torch using torch.multiprocessing.set_start_method(‘spawn’)

Only torch.Tensor types will be considered