From 15204d569b26dd78a1cd0afee9317e378849493c Mon Sep 17 00:00:00 2001 From: Abdul Fatir Ansari Date: Sat, 27 May 2023 10:29:21 +0200 Subject: [PATCH] Fix MQF tests --- src/gluonts/torch/model/mqf2/lightning_module.py | 2 ++ test/torch/model/test_mqf2_modules.py | 1 + 2 files changed, 3 insertions(+) diff --git a/src/gluonts/torch/model/mqf2/lightning_module.py b/src/gluonts/torch/model/mqf2/lightning_module.py index 6dc824beb4..470eee8d58 100644 --- 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 index 85fa21337f..451a16d890 100644 --- 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]