Calculating affine transformation matrix

So I am trying to learn PyTorch and as an experiment I tried to apply a specific geometric transform (rotation by 45 degrees) to an image using torch.nn.functional.affine_grid and torch.nn.functional.grid_sample. And now from this new image I try to work backwards using the MSE between the images to recover the transformation matrix. The loss does seem to come down but ultimately it doesn’t quite do the job.

Here’s a link to a gist with the notebook: https://gist.github.com/abnan/19782d46e8a8396267d2aee44b48a3de

Any pointers on what I am doing wrong? Also in general what/how to debug such cases?

if you want to go more high level you can use Kornia library
https://kornia.readthedocs.io/en/latest/geometry.transform.html#kornia.geometry.transform.rotate