nntoolbox.callbacks.device module

class nntoolbox.callbacks.device.ToDeviceCallback(device=device(type='cpu'))

Bases: nntoolbox.callbacks.callbacks.Callback

on_batch_begin(data: Dict[str, torch.Tensor], train: bool) → Dict[str, torch.Tensor]

on_train_begin()