GPU Acceleration

GPU training is supported, but you must install the backend package for your hardware first.

NVIDIA GPUs

using Pkg
Pkg.add("LuxCUDA")
# or
Pkg.add(["CUDA", "cuDNN"])

AMD ROCm GPUs

using Pkg
Pkg.add("AMDGPU")
# If your use case fails, consider Reactant.jl for GPU support.

Metal M-Series GPUs

using Pkg
Pkg.add("Metal")

Intel GPUs

using Pkg
Pkg.add("oneAPI")

Then run the following to access a device:

NVIDIA GPUs

using EasyHybrid, LuxCUDA
gpu_device()

AMD ROCm GPUs

using EasyHybrid, AMDGPU
gpu_device()

Metal M-Series GPUs

using EasyHybrid, Metal
gpu_device()

Intel GPUs

using EasyHybrid, oneAPI
gpu_device()

In your training call, pass arch = GPU(). For example:

using EasyHybrid, Metal

train(...; arch = GPU())

That is all you need. Your hybrid model will now train on the GPU.