Not resetting state of charge

src/trainers/quantile_trainer.py

@@ -557,6 +557,9 @@ class NonAutoRegressiveQuantileRegression(Trainer):
                 inputs, targets = inputs.to(self.device), targets.to(self.device)
 
                 outputs = self.model(inputs)
+
+                outputs = outputs.reshape(-1, len(self.quantiles))
+
                 outputted_samples = [
                     sample_from_dist(self.quantiles, output.cpu()) for output in outputs
                 ]