So I thought I’d simplify things a little and I think I’ve hit the problem. I suspect this is something to do with the dynamic graph nature of PyTorch. The following code comes up with a reasonable result:
import torch

use_cuda = True
device = torch.device("cuda" if use_cuda else "cpu")

x_rot = torch.tensor([0.0], dtype=torch.float32,\
  requires_grad=True, device = device)

x_sin = torch.sin(x_rot)
x_cos = torch.cos(x_rot)

x_sin_mask = torch.tensor([[0,0,0,0],\
    [0,0,-1,0],
    [0,1,0,0],
    [0,0,0,0]], dtype=torch.float32, device=device)

x_cos_mask = torch.tensor([[0,0,0,0],\
    [0,1,0,0],
    [0,0,1,0],
    [0,0,0,0]], dtype=torch.float32, device=device)

base = torch.tensor([[1,0,0,0],\
    [0,0,0,0],\
    [0,0,0,0],\
    [0,0,0,1]], dtype=torch.float32, device=device)

rot_x = x_cos.expand_as(x_cos_mask) * x_cos_mask +\
    x_sin.expand_as(x_sin_mask) * x_sin_mask + base

print(rot_x)
learning_rate = 0.1

# Base is rotated by 90 around X
base = torch.tensor([[1,0,0,0],
  [0,0,-1,0],
  [0,1,0,0],
  [0,0,0,1]], dtype=torch.float32, device=device)

loss = ((rot_x-base)**2).sum()
loss.backward()

with torch.no_grad():
  x_rot -= learning_rate * x_rot.grad 
  x_rot.grad.zero_()
  print(x_rot)


However, if I attempt to put a loop around the following lines:

for i in range(20):
  loss = ((rot_x-base)**2).sum()
  loss.backward()

  with torch.no_grad():
    x_rot -= learning_rate * x_rot.grad 
    x_rot.grad.zero_()
    print(x_rot)


We get the problem where no gradient appears on x_rot. I suspect I need to find out more about loops and the backing graph that pytorch creates.