diff --git a/tests/imputation/timemixer.py b/tests/imputation/timemixer.py
index bf5c72cd..a0735663 100644
--- a/tests/imputation/timemixer.py
+++ b/tests/imputation/timemixer.py
@@ -47,8 +47,8 @@ class TestTimeMixer(unittest.TestCase):
         DATA["n_features"],
         n_layers=2,
         top_k=5,
-        d_model=512,
-        d_ffn=512,
+        d_model=32,
+        d_ffn=32,
         dropout=0.1,
         epochs=EPOCHS,
         saving_path=saving_path,