offload_to_cpu

Offload Module to CPU

Description

Moves a torch module and all its parameters to CPU.

Usage

offload_to_cpu(module, gc = TRUE)

Arguments

  • module: A torch nn_module.
  • gc: Logical. Run garbage collection after offload.

Value

The module (modified in place).

Examples

model$to(device = "cuda")
output <- model(x)
offload_to_cpu(model)