Could you post the shapes of log_p and target?
Note that your target should also contain the batch dimension ([n, h, w]) and should contain class indices in the range [0, nb_classes-1] in the usual use case.
It looks like you are selecting all “valid” pixels using the mask and store them as a flattened version. I would assume that you have to select the corresponding targets in the same way so that the shapes match again.
However, you could probably also use the original output and target with a non-reducing criterion (pass reduction='none' tonn.NLLLoss`) and multiply it with your mask to zero out the unwanted pixel locations.
We could most likely avoid the for loop by using broadcasting, so could you also post the shape of your mask?