将Pytorch设置为在AMD GPU上运行。

Question

4

根据官方文档，PyTorch现在支持AMD GPU。ROCm 4.2可以通过pip安装。但我无法在Google或官方文档中找到如何强制DL训练使用GPU的方法。以下命令的AMD等效命令是什么？

torch.device('cuda' if torch.cuda.is_available() else 'cpu')

- castrovictor

2

Reference: https://dev59.com/zMDqa4cB1Zd3GeqPWyJM - user14467343

1个回答

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- JustOneMan · Accepted Answer

ROCm可以转化CUDA代码，因此torch.cuda.is_available()应该可以正常工作。