Search code examples
juliaflux.jl

How to use Flux.jl on an AMD GPU?


I currently do not have an Nvidia GPU on my laptop but I still want to be able to speed up training. Per the Flux docs, it says Nvidia GPU's are supported out of the box but doesn't mention AMD GPU's at all. Is it possible to work with Flux on an AMD GPU?


Solution

  • This should in principle be possible, since AMDGPU.jl provides a similar interface for AMD GPUs as CUDA.jl does for NVIDIA GPUs, and Flux is claimed to be agnostic to array types. However, at best this will only work on Linux, since AMDGPU.jl relies on AMD's ROCm platform, and ROCm is only supported on Linux, since ROCm is indeed specifically built around the Linux kernel.