Small refactor of IoUs computations to be sure that only one way of logging them is used (cf. Common Pitfalls of using metrics in lightning). After updating to torch 2.1 we would have one way to log mean IoU and another way to log class IoUs. No anomaly was observed, but better safe than sorry.
Mixing the two logging methods by calling self.log("val", self.metric) in {training|validation|test}step method and then calling self.log("val", self.metric.compute()) in the corresponding on{train|validation|test}_epoch_end method. Because the object is logged in the first case, Lightning will reset the metric before calling the second line leading to errors or nonsense results.