Training with GPU

To use a GPU, set the backend under TrainConfig to "gpu" in your YAML config. Nothing else in your training code changes.

TrainConfig:
  # ...
  backend: "gpu"                    # compute backend: "cpu" or "gpu"
  # ...

Requirements

• NVIDIA GPU with recent drivers installed
• CUDA-enabled environment configured (and available to your Python stack)
• If you run into “CUDA not available” errors, switch backend back to "cpu" or fix your CUDA setup

# Quick check that the GPU is visible
nvidia-smi