嘿，这个问题我刚好折腾过！在ROCm环境下，对应CUDA里CUDA_VISIBLE_DEVICES的环境变量是ROCR_VISIBLE_DEVICES，用法和CUDA的那个几乎一模一样，上手毫无门槛。

下面两种常见的设置方式，你可以按需选：

• 终端运行前临时设置：直接在启动Python脚本的命令前加上环境变量指定，比如只想用第0块GPU，就这么写：

  ROCR_VISIBLE_DEVICES=0 python your_training_script.py
  

  如果要选多块（比如第1和第3块），就用逗号分隔：ROCR_VISIBLE_DEVICES=1,3 python your_script.py；要是想禁用所有GPU，写ROCR_VISIBLE_DEVICES=-1就行。

• 在Python脚本内部设置：和你用os.environ设置CUDA变量的逻辑完全一致，但有个关键前提——必须在导入任何ROCm相关的库（比如PyTorch、TensorFlow）之前设置，不然不会生效。举个实际例子：

  import os
  # 先搞定环境变量，再碰其他库！
  os.environ["ROCR_VISIBLE_DEVICES"] = "0,2"
  # 现在才导入依赖
  import torch
  # 验证一下是否生效
  print(f"可用GPU数量: {torch.cuda.device_count()}")  # 这里应该输出2
  

另外提个小细节：设备编号是从0开始计数的，和CUDA的规则完全一致，不用额外转换，直接用就行。

备注：内容来源于stack exchange，提问作者mmcmp