Update test

amazon-science · Mar 25, 2024 · 77b5244 · 77b5244
1 parent 9c44b7a
commit 77b5244
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/test/test_chronos.py b/test/test_chronos.py
@@ -186,12 +186,9 @@ def test_pipeline_embed(torch_dtype: str):
         device_map="cpu",
         torch_dtype=torch_dtype,
     )
-    model_context_length = pipeline.model.config.context_length
-    expected_embed_length = model_context_length + (
-        1 if pipeline.model.config.use_eos_token else 0
-    )
     d_model = pipeline.model.model.config.d_model
     context = 10 * torch.rand(size=(4, 16)) + 10
+    expected_embed_length = 16 + (1 if pipeline.model.config.use_eos_token else 0)
 
     # input: tensor of shape (batch_size, context_length)