Output the class for each grid cell of an yolo-detector (probability map))

Is it possible to achieve a similar output to the probability map shown in the yolo paper(see below)?

I have already seen in numerous documents on the yolov1 and v2 cnn that a probability map is often presented. For my application, I would like to output just such a map for the detection of an image instead of the bounding box. That means I want to have a unique assignment to one of x classes for each cell of my grid.

So far, however, I have not been able to make such an output possible.

I multiplied the probability with all class scores and searched for the maximum in all anchor boxes. This leads to the fact that in many cells, the wrong classes are found. The responsible cells on which the center of the detected boxes falls are displayed correctly with a high score.