revert changes to Model

merlin/models/torch/models/base.py

@@ -88,16 +88,19 @@ def forward(
     def training_step(self, batch, batch_idx):
         """Performs a training step with a single batch."""
         del batch_idx
-        loss_and_metrics = self._step(batch)
+        if isinstance(batch, Batch):
+            features = batch.features
+            targets = batch.targets
+        else:
+            features, targets = batch
 
-        return loss_and_metrics["loss"]
+        predictions = self(features, batch=Batch(features, targets))
 
-    def validation_step(self, batch, batch_idx):
-        """Performs a validation step with a single batch."""
-        del batch_idx
-        loss_and_metrics = self._step(batch)
+        loss_and_metrics = compute_loss(predictions, targets, self.model_outputs())
+        for name, value in loss_and_metrics.items():
+            self.log(f"train_{name}", value)
 
-        return loss_and_metrics
+        return loss_and_metrics["loss"]
 
     def configure_optimizers(self):
         """Configures the optimizer for the model."""
@@ -115,21 +118,6 @@ def last(self) -> nn.Module:
         """Returns the last block in the model."""
         return self.values[-1]
 
-    def _step(self, batch):
-        if isinstance(batch, Batch):
-            features = batch.features
-            targets = batch.targets
-        else:
-            features, targets = batch
-        predictions = self(features, batch=Batch(features, targets))
-
-        loss_and_metrics = compute_loss(predictions, targets, self.model_outputs())
-        train_or_eval = 'train' if self.training else 'valid'
-        for name, value in loss_and_metrics.items():
-            self.log(f"{train_or_eval}_{name}", value)
-
-        return loss_and_metrics
-
 
 def compute_loss(
     predictions: Union[torch.Tensor, Dict[str, torch.Tensor]],