Fix MQF tests

awslabs · May 27, 2023 · 15204d5 · 15204d5
1 parent 8f83698
commit 15204d5
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/src/gluonts/torch/model/mqf2/lightning_module.py b/src/gluonts/torch/model/mqf2/lightning_module.py
@@ -96,6 +96,7 @@ def _compute_loss(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor:
         future_time_feat = batch["future_time_feat"]
         future_target = batch["future_target"]
         past_observed_values = batch["past_observed_values"]
+        future_observed_values = batch["future_observed_values"]
 
         picnn = self.model.picnn
 
@@ -107,6 +108,7 @@ def _compute_loss(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor:
             past_observed_values,
             future_time_feat,
             future_target,
+            future_observed_values,
         )
 
         hidden_state = hidden_state[:, : self.model.context_length]

diff --git a/test/torch/model/test_mqf2_modules.py b/test/torch/model/test_mqf2_modules.py
@@ -78,6 +78,7 @@ def test_mqf2_modules(
         past_observed_values,
         future_time_feat,
         future_target,
+        future_observed_values,
     )
 
     hidden_state = hidden_state[:, :context_length]