If your goal is to train a network that find image transformation, you can create a model that predicts the transformation (ie 2x3 matrix for fulls 2d affine, or predict 3 parameters rotationZ, translationX, translationY and generate the 2x3 matrix by yourself in order to ensure “rigid” transformation).
Then apply this transformation to the original image and compute loss.
See