Currently, we cannot compute train metrics because it requires to save all the preds and targets over the epoch.
The problem also affects the val and test sets, but since they're smaller, the effect is less noticeable.
To fix this, the MetricWrapper should have an update method taking the preds and target that calls the underlying self.metrics.update, and the compute method which no longer takes in the pred and target, but instead calls the self.metric.compute
Also, all TorchMetrics in spaces.py should become their class equivalent rather than functions.
Currently, we cannot compute train metrics because it requires to save all the preds and targets over the epoch.
The problem also affects the val and test sets, but since they're smaller, the effect is less noticeable.
To fix this, the MetricWrapper should have an update method taking the preds and target that calls the underlying self.metrics.update, and the compute method which no longer takes in the pred and target, but instead calls the self.metric.compute
updateto theMetricWrapperMetricWrapper.computeto work with the updateAlso, all TorchMetrics in spaces.py should become their class equivalent rather than functions.
spaces.pyto use classes rather than functions. Make sure the classes get initialized.