To solve the problem with retain_graph, can you try cloning and detaching rot_mat in the beginning of for loop? i.e., I feel that you have a long chain of rot_map in memory which makes it difficult to trace. Can you try this inside for loop?
for i in range(100):
    rot_mat = rot_mat.detach().clone()
    rot_mat.requires_grad=True
    ....