Fix missing checkpoint for forecast engine (#1771)

* implemented

* remove eval in interface

* lint

* incoporate requested changes

* fix imports

* Fixed missing checkpoint

---------

Co-authored-by: moritzhauschulz <[email protected]>

Author: clessig

src/weathergen/model/model.py

@@ -18,6 +18,7 @@
 import numpy as np
 import torch
 import torch.nn as nn
+from torch.utils.checkpoint import checkpoint
 
 from weathergen.common.config import Config
 from weathergen.datasets.batch import ModelBatch
@@ -585,7 +586,7 @@ def forward(self, model_params: ModelParams, batch: ModelBatch) -> ModelOutput:
         for step in batch.get_output_idxs():
             # apply forecasting engine (if present)
             if self.forecast_engine:
-                tokens = self.forecast_engine(tokens, step)
+                tokens = checkpoint(self.forecast_engine, tokens, step)
 
             # decoder predictions
             output = self.predict_decoders(model_params, step, tokens, batch, output)