Update test

awslabs · May 26, 2023 · 0aa8e60 · 0aa8e60
1 parent 67947f4
commit 0aa8e60
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/test/torch/model/test_deepar_modules.py b/test/torch/model/test_deepar_modules.py
@@ -79,6 +79,7 @@ def test_deepar_modules(
         past_observed_values,
         future_time_feat,
         future_target,
+        future_observed_values,
     )
 
     assert scale.shape == (batch_size, 1)
@@ -231,6 +232,11 @@ def test_rnn_input(
         dtype=torch.float32,
     ).view(1, prediction_length)
 
+    batch["future_observed_values"] = torch.ones(
+        (1, prediction_length),
+        dtype=torch.float32,
+    )
+
     rnn_input, scale, _ = model.prepare_rnn_input(**batch)
 
     assert (scale == 1.0).all()