Got it, thanks!
I installed Pytorch from source and now I have access to the grad and backward functions in torch.autograd.
However when I try to use grad my kernel just crashes. I use like this:
output_img = model.forward(input_img)
g = grad(output_img, input_img, create_graph=True)
Is this not intended to be used with non-scalar (multiple-dimensional) Tensors? The examples on the github work fine.