So I think I’ve solved it. It’s a little verbose but this loss function appears to work for rotation matrices. I figure I’ll post this here in case it’s useful to folks?
Thanks for the help ya’all!
""" A short program to test the loss function for a rotation matrix.
We take in 3 angles and attempt to move one matrix to another.
We default to the GPU.
"""

import torch
import math

def gpu(myfunc):
  def inner_func(*args, **kwargs):
    device = torch.device("cuda")
    tt = myfunc(*args, **kwargs)
    tt.to(device)
    return tt
  return inner_func

learning_rate = 0.1
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)

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

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

def make_rot(x_rot, y_rot, z_rot, device="cpu") : 
  """ Make a rotation matrix from 3 tensors of dimension [1]
  representing the angle in radians around X, Y and Z axes in this
  order. It seems very verbose but this really does seem to work.
  """
  x_sin = torch.sin(x_rot)
  x_cos = torch.cos(x_rot)

  y_sin = torch.sin(y_rot)
  y_cos = torch.cos(y_rot)

  z_sin = torch.sin(z_rot)
  z_cos = torch.cos(z_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)

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

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

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

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

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

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

  base_z = torch.tensor([[0,0,0,0],\
      [0,0,0,0],\
      [0,0,1,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_x

  rot_y = y_cos.expand_as(y_cos_mask) * y_cos_mask +\
      y_sin.expand_as(y_sin_mask) * y_sin_mask + base_y

  rot_z = z_cos.expand_as(z_cos_mask) * z_cos_mask +\
      z_sin.expand_as(z_sin_mask) * z_sin_mask + base_z

  # Why does this line screw up but two matmuls are fine? :S
  #rot_mat = rot_x * rot_y * rot_z
  tmat = torch.matmul(rot_x, rot_y)
  rot_mat = torch.matmul(tmat, rot_z)
  return rot_mat 

if __name__ == "__main__":

  base = make_rot(\
    torch.tensor([math.radians(90)], dtype=torch.float32),\
    torch.tensor([math.radians(45)], dtype=torch.float32),\
    torch.tensor([math.radians(10)], dtype=torch.float32),\
    device = device)

  print(base)

  for i in range(50):
    print("Step", i)
    rot_mat = make_rot(x_rot, y_rot, z_rot, device)
    print(rot_mat)
    loss = ((rot_mat-base)**2).sum()
    loss.backward()
    print(loss)

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

      print(x_rot, y_rot, z_rot)