安装方法:
Enable PyTorch with DirectML on Windows | Microsoft Learn
新的pytorch-directml 不再是独立的pytorch 移植,变成了一个设备插件,更好用,兼容性更好。
原来的版本无法跑transformers, 新版变成一个独立的计算设备 dml, 兼容性更好
import torch
import torch_directml
dml=torch_directml.device()
tensor1 = torch.tensor([1]).to(dml) # Note that dml is a variable, not a string!
tensor2 = torch.tensor([2]).to(dml)
for i in range(1000000):
dml_algebra = tensor1 + tensor2
dml_algebra.item()