我有一个a
由神经模型生成的向量,需要与一个巨大的矩阵交互M
。由于M
很大,我必须在 cpu 设备中进行计算。在这种情况下,我想知道梯度是否可以在cuda设备上保留并向后移动。
示例如下:
a_cuda = torch.randn([1, 512], requires_grad=True).to("cuda")
a_cpu = torch.randn([1, 512], requires_grad=True).to("cpu")
M = torch.randn([512, 100000], requires_grad=False) # loaded in cpu device, dont need update
out_cuda = (a_cuda.cpu() @ M).sum()
out_cuda.backward()
out_cpu = (a_cpu @ M).sum()
out_cpu.backward()
print(a_cuda.grad) # None
print(a_cpu.grad)
我正在寻找具有a_cuda.grad
相同梯度的解决方案a_cpu.grad
。