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: