Skip to content

[utils]

Utils for torch backend support code.

GPU/CPU backing device management utils

  • AutoDeviceModule: Wrapper for a torch.nn.Module, automating device-management.
  • select_device: Select a backing device to use based on inputs and availability.