CODEWITHSUNDEEP
pytorch
r4bb1t
r4bb1t

Reputation: 1165

How do I retain grads and also change device type in pytorch?

When I change my input variable from CPU to cuda, it loses all its grad and also loses its is_leaf status. How do I circumvent this? I want to keep the gradients and also change it to another device.

Upvotes: 1

Views: 927

Answers (1)

jodag
jodag

Reputation: 22204

A leaf tensor is one which has its requires_grad attribute set to True. When you do any out-of-place operation on a tensor the resulting tensor is no longer a leaf tensor. This includes creating a copy of the tensor on a different device using .to(device), .cuda(), or .cpu(). The recommended way to set the requires_grad attribute of an existing tensor is to use the in-place Tensor.requires_grad_ method. If you want the tensor on the GPU to be the leaf node then you will need to set requires_grad after copying to the desired device.

For example

input = input.to('cuda')
input.requires_grad_(True)   # need to set requires_grad after copying to GPU

or a little more concisely

input = input.to('cuda').requires_grad_(True)

Upvotes: 2

Related Questions