Thanks, @SimonW. Below is a MWE that is much simpler than what I had posted earlier. Turns out that the deadlock is caused only when N=1; values greater than 1 don’t seem to cause the deadlock.
import torch
from torch.autograd import Variable

N = 1   # N = 1 CAUSES DEADLOCK, N > 1 DOESN'T
D = 2

x = Variable(torch.randn(N, D), requires_grad=False)
y = Variable(torch.randn(N, D), requires_grad=False)

w1 = Variable(torch.randn(D,D), requires_grad=True)

y_pred = x.mm(w1)

loss = (y_pred - y).pow(2).sum()

learning_rate = 1e-6
print('Check1')
loss.backward(retain_graph=True)
print('Check2')

w1.data -= learning_rate * w1.grad.data

w1.grad.data.zero_()