modify output

MannLabs · May 24, 2024 · fc4f256 · fc4f256
1 parent 16df970
commit fc4f256
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/sparcscore/ml/plmodels.py b/src/sparcscore/ml/plmodels.py
@@ -213,7 +213,7 @@ def training_step(self, batch):
         self.log('mse/train', self.mse(output, target.unsqueeze(1)), on_step=False, on_epoch=True, prog_bar=True)
         self.log('mae/train', self.mae(output, target.unsqueeze(1)), on_step=False, on_epoch=True, prog_bar=True)
 
-        return loss
+        return {'loss': loss, 'predictions': output, 'targets': target}
 
     def validation_step(self, batch):
         data, target = batch
@@ -226,7 +226,7 @@ def validation_step(self, batch):
         self.log('mse/val', self.mse(output, target.unsqueeze(1)), on_step=False, on_epoch=True, prog_bar=True)
         self.log('mae/val', self.mae(output, target.unsqueeze(1)), on_step=False, on_epoch=True, prog_bar=True)
 
-        return loss
+        return {'loss': loss, 'predictions': output, 'targets': target}
 
     def test_step(self, batch):
         data, target = batch
@@ -239,7 +239,7 @@ def test_step(self, batch):
         self.log('mse/test', self.mse(output, target.unsqueeze(1)), on_step=False, on_epoch=True, prog_bar=True)
         self.log('mae/test', self.mae(output, target.unsqueeze(1)), on_step=False, on_epoch=True, prog_bar=True)
 
-        return loss
+        return {'loss': loss, 'predictions': output, 'targets': target}
 
 # implemented models for future use currently not applied to SPARCSpy