Hi all, I’m making a network which predict objects in output grids so that the output shape is [class_size, grid_width, grid_height]. Then now the loss function is like this. output is [batch_size, class_size, grid_width, grid_height] and  target [1, grid_width, grid_width] contains class labels if there is an object in a grid otherwise 0.
for b, w, h in torch.nonzero(target.data):
    loss += F.cross_entropy(output[b, :, w, h].contiguous().view(1, -1), target[b, w, h])

As you may think, this loss function is very slow and the bottleneck so that I’d like to change if possible.
The reason why the target is tensor is I use the DataLoader and it does not accept some objects, for example. list which store grid and corresponding class label like list(dict(grid=(x,y), cls=c). dict(grid=(x,y), cls=c)...) for each image.
Does anyone have good idea for improvement the performance? Thank you for advance.