Incorrect Z-axis rotation matrix in `get_random_rot` method

[gregspangenberg]

Hi, thank you for creating this i have found it quite useful. I just wanted to bring this to your attention.

Issue description

There's an incorrect implementation of the Z-axis rotation matrix in the get_random_rot method. The current code incorrectly implements a rotation around the X-axis instead of the Z-axis.

Current implementation (incorrect)

R3 = torch.eye(3)
R3[1, 1] = torch.cos(theta_z)
R3[2, 2] = torch.cos(theta_z)
R3[1, 2] = -1*torch.sin(theta_z)
R3[2, 1] = torch.sin(theta_z)

Correct implementation

R3 = torch.eye(3)
R3[0, 0] = torch.cos(theta_z)
R3[1, 1] = torch.cos(theta_z)
R3[0, 1] = -1*torch.sin(theta_z)
R3[1, 0] = torch.sin(theta_z)

This bug affects the model's ability to properly learn rotation equivariance during training.