需要先把Tensor放入到GPU中,然后再设置Tensor.requires_grad=True。
from torch import optim
weights = torch.rand(2,1,128,416)
weights.requires_grad = True
weights = weights.cuda()
optimizer = optim.Adam([weights], lr=0.01)
.
.
.
这样写会报“can't optimize a non-leaf Tensor”的错。
from torch import optim
weights = torch.rand(2,1,128,416)
weights = weights.cuda()
weights.requires_grad = True
optimizer = optim.Adam([weights], lr=0.01)
.
.
.
这样写就好了。