Add separate EPD that trains on processor

Author: sgreenbury

src/auto_cast/models/encoder_processor_decoder.py

@@ -134,3 +134,40 @@ def _advance_batch(self, batch: Batch, next_inputs: Tensor, stride: int) -> Batc
             constant_scalars=batch.constant_scalars,
             constant_fields=batch.constant_fields,
         )
+
+
+class EPDTrainProcessor(EncoderProcessorDecoder):
+    """Encoder-Processor-Decoder Model training on processor."""
+
+    train_processor: Processor
+
+    def __init__(
+        self,
+        encoder_decoder: EncoderDecoder,
+        processor: Processor,
+        learning_rate: float = 1e-3,
+        stride: int = 1,
+        teacher_forcing_ratio: float = 0.5,
+        max_rollout_steps: int = 10,
+        loss_func: nn.Module | None = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(
+            encoder_decoder=encoder_decoder,
+            processor=processor,
+            learning_rate=learning_rate,
+            stride=stride,
+            teacher_forcing_ratio=teacher_forcing_ratio,
+            max_rollout_steps=max_rollout_steps,
+            loss_func=loss_func,
+            **kwargs,
+        )
+
+    def training_step(self, batch: Batch, batch_idx: int) -> Tensor:  # noqa: ARG002
+        encoded_batch = self.encoder_decoder.encoder.encode_batch(batch)
+        output = self.processor.map(encoded_batch.encoded_inputs)
+        loss = self.processor.loss(output, encoded_batch.encoded_output_fields)
+        self.log(
+            "train_loss", loss, prog_bar=True, batch_size=batch.input_fields.shape[0]
+        )
+        return loss