Make sure to add it to all_results here, as follows:
all_results = torch.cat(
[
torch.unsqueeze(sum_result, dim=1),
torch.unsqueeze(average_result, dim=1),
torch.unsqueeze(count_result, dim=1),
torch.unsqueeze(portion_result, dim=1)
],
dim=1,
)
This is the only thing you need to do, it will be added to the loss automatically. To compute portion_result, you can do:
portion_result = count_result / some dimension of scaled_probability_per_cell
Can you print out the shape of scaled_probability_per_cell?