Multiple output together with RocAuc(or any Accumelator)

Hey everyone.
I was working with a network that output a dict of outputs, e.g {‘x’:…, ‘z’:…, ‘t’:…}
Adding a metric to such network is easy enough, just adding a special version of the function which first extract the arguments.
But for case like RocAuc one need to first aggregate those outputs, and only at the end of the epoch run those computation.
But overriding the AUC is a bit tricky, as it is by itself not a function, but an Accumulator.

Here is an example for simple metric overloading.

How would you do this?

Ido