Differentiable affine transformation on patches of images in pytorch

I have a tensor of object bounding boxes, e.g. with the shape of [10,4] which correspond to a batch of images e.g. with shape [2,3,64,64] and transformation matrices for each object with shape [10,6] and a vector that defines which object index belongs to which image. I would like to apply the affine transformations on patches of the images and replace those patches after applying the transformations. I am doing this using a for loop now, but the way I am doing it is not differntiable (I get the in place operation error from pytorch). I wanted to know if there is a differntiable way to do this. e.g. via grid_sample?

Here is my current code:

for obj_num in range(obj_vecs.shape[0]): #batch_size
    im_id = obj_to_img[obj_num]
    x1, y1, x2, y2 = boxes_pred[obj_num]
    im_patch = img[im_id, :, x1:x2, y1:y2]
    im_patch = im_patch[None, :, :, :]
    img[im_id, :, x1:x2, y1:y2] = self.VITAE.stn(im_patch, theta_mean[obj_num], inverse=False)[0]

using im_batch.clone() instead of feeding it directly to STN solved the issue.