Log and save model parameter counts

test_traj_mixer.py

@@ -8,6 +8,7 @@ from models import (
     validate_traj_mixer_config,
     validate_traj_mixer_state_dict,
 )
+from train_util import get_model_parameter_counts
 
 
 class TrajMixerTest(unittest.TestCase):
@@ -91,6 +92,16 @@ class TrajMixerTest(unittest.TestCase):
         with self.assertRaisesRegex(ValueError, "divisible"):
             TrajMixer(n_embd=121, n_head=10, hidden_group=20)
 
+    def test_parameter_counts_match_traj_mixer_parameters(self) -> None:
+        mixer = TrajMixer(n_embd=120, n_head=10, hidden_group=20)
+        self.assertEqual(
+            get_model_parameter_counts(mixer),
+            {
+                "model_parameter_count": 8_640,
+                "trainable_parameter_count": 8_640,
+            },
+        )
+
 
 if __name__ == "__main__":
     unittest.main()