How to transform affine matrix using affine transformation

I have an affine transformation matrix to transform a 3x3x3 MR image. I have transformed the image using PyTorch and created the matrix for the new image. But to save the image in nii.gz format, I need to compute the affine matrix for the transformed image. I could not find any resources regarding how to do that. I tried applying the transformation on the affine matrix of the image but that didn’t work. Any help is highly appreciated.

Thanks,