I have a bunch of computations like this:
start_coords has shape [b,2,2] with start and end ranges for each image
img_available = torch.zeros_like(t)
img_available[:,start_coords[:,0,0]:start_coords[:,0,1],start_coords[:,1,0]:start_coords[:,1,1]]=1

or the other way around:
selected = imgs[:, start_coords[:,0,0]:start_coords[:,0,1],start_coords[:,1,0]:start_coords[:,1,1]]

so for each image i in batch, i need imgs[i, start_coords[i,0,0]:start_coords[i,0,1],start_coords[i,1,0]:start_coords[i,1,1]]
but this is not working. I need some way to assign/get values using batched-ranges. The length is always the same.
This is inside a loop, so it needs to reasonably fast!