Workaround to indexing containing grads

Hello community,

I am struggling to think of a way to get the behavior I am looking for.

I have a sub-network predicting 3D regions of interest, yielding a tensor of the shape (num_objects, 6). The 6 refers to x1, y1, z1, x2, y2, z2 coordinates of the 3D region of interest. Based on these coordinates I want index a boolean tensor that has the same shape as the input image, namely (120, 120, 120), and set the values of these region of interests to False.
So naturally I first cast the tensor of shape (num_objects, 6) to integers and use the predicted coordinates for the sake of indexing.

dummy_fmap[roi[0][0]:roi[0][3], roi[0][1]:roi[0][4], roi[0][2]:roi[0][5]] = False

Since the indexing operation is unfortunately not differentiable, the weights in the sub-network are not optimized. Now I am thinking of alternate approaches and would be glad to get some help on this.
I am aware that ROI pooling of two stage detectors does something similar, but was not able to adopt this to my task.

Thanks in advance!