Modifying the prediction function in training loop

Part 1 (2020)
1

If my prediction function is not calculated by taking argmax. How can I make changes in calculating metrics while training?
Accumulator class
How do I make changes in the accumulator class and add it in the training loop?