Adding aggregation to TAPAS

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?