Setting up Pytorch in Venv for both CUDA and ROCm

Question

I want to install PyTorch on multiple different machines with different hardwares, using both CUDA and ROCm hardware, but I'm not sure how to setup a venv that works for both, so to have as simple an environment solution as possible for all.

I've only searched online to see if there's anyone that has had similar problems before. I couldn0t find anything

Answer 1

I suggest to make:

On the CUDA machine:
1 virtual environment
1 file requirements_CUDA.txt with the related pytorch version
install with:
pip install -r requirements_CUDA.txt

On the ROCm machine 1 virtual environment
1 file requirements_ROC.txt with the related pytorch version
install with:
pip install -r requirements_ROC.txt i