model_to#
tiatoolbox.utils.misc.model_to
- model_to(model, *, on_gpu)[source]#
Transfers model to cpu/gpu.
- Parameters:
model (torch.nn.Module) – PyTorch defined model.
on_gpu (bool) – Transfers model to gpu if True otherwise to cpu.
- Returns:
The model after being moved to cpu/gpu.
- Return type: