Skip to content

PyTorch on ROCm FAQ

Jump to bottom
Jithun Nair edited this page Aug 3, 2023 · 1 revision

Q. I'm using an environment with ROCm installed (e.g. in /opt/rocm) and PyTorch ROCm wheels installed. I'm seeing errors when doing import torch e.g.

terminate called after throwing an instance of 'std::bad_alloc'
  what():  std::bad_alloc
Aborted (core dumped)

A. The PyTorch wheels with ROCm support already bundle the ROCm libraries needed for PyTorch. These libraries (.so files) are installed in the same path wherever PyTorch is installed e.g. /usr/local/lib/python3.8/dist-packages/torch/lib (depends on your torch package installation path). However, if you have LD_LIBRARY_PATH environment variable defined to point to the lib path for ROCm installation e.g. opt/rocm/lib (depends on your ROCm installation path), that results in the loader picking up potentially incompatible ROCm libraries from there instead of the ones bundled with the wheel.

To resolve this, add the torch package's lib path in front of the ROCm lib path in your LD_LIBRARY_PATH e.g.
export LD_LIBRARY_PATH="/usr/local/lib/python3.8/dist-packages/torch/lib:$LD_LIBRARY_PATH"

Clone this wiki locally