Get tensor value at some coordinate along batch with loader giving out dicts

my dataloader produces dictionaries as outputs, which are input for a neural net.

let’s say:

def dummy_get_item():
image = torch.rand((1,64,64))
inpainted_image = image.clone()
mask_locs = torch.tensor([5,10,15,22]) #randomly chosen
inpainted_image[0,mask_locs[0]:mask_locs[1],mask_locs[2]:mask_locs[3]] = 0
return {'inpainted_image':inpainted_image,'truth_image':image,'locs':mask_locs}

i want to be able to measure loss for only the values i silenced, but it’s tricky because from each sample in the batch i’m trying to collect different coordinates
say batch_size = 4

input_dict['inpainted_image'].shape = [4,1,64,64]
input_dict['mask_locs'].shape = [4,4]

i try:

ind = input_dict['mask_loss']
loss = criterion(input_dict['truth_image'][:,ind[:,0]:ind[:,1],ind[:,2]:ind[:,3]],input_dict['inpainted_image'][:,ind[:,0]:ind[:,1],ind[:,2]:ind[:,3]])

what i’m trying to do is compare the values of “truth image” and “inpainted image” only at the locations where it was masked.

i’m getting the following error:
TypeError: only integer tensors of a single element can be converted to an index


update - i figured out i can solve this by doing
locs = (inpainted_image==0) * truth_image!=0)
and then solve the batch issue.

but i still would love to know how can i solve it in the logic i presented before, meaning how to overcome the error i got.