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: