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