CommonLoopUtils
CommonLoopUtils copied to clipboard
Makes Metric/Collection inherit from PyNodeTree.
Makes Metric/Collection inherit from PyNodeTree.
That avoids having to decorate Metric/Collection subclasses with @flax.struct.dataclass and removes one class of confusing errors.
Note that all calling code has to be updated accordingly, because subclasses of PyNodeTree can NOT be decorated with @flax.struct.dataclass anymore.