In case anybody is interested. I managed to obtain the desired results by adding scaling to theta according to the image size:
    theta = torch.cat((rot_mat(parameters[:, 0]), parameters[:, 1:, None]), dim=-1)
    theta[:, :, 0] *= shape[1] / 2
    theta[:, :, 1] *= shape[0] / 2

    scale = 1
    n, c, rows, cols = orig_image.shape
    sample_grid = F.affine_grid(theta, (n_images, 1, int(rows * scale), int(cols * scale)), align_corners=False).to(orig_image.device)
    sample_grid[..., 0] /= shape[1] / 2
    sample_grid[..., 1] /= shape[0] / 2

    images = F.grid_sample(orig_image, sample_grid, align_corners=False)

giving the following result