Scale TrajMixer hidden width with head count

models.py

@@ -15,7 +15,7 @@ from backbones import (
 from targets import PAD_IDX
 
 
-TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v1"
+TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v2"
 
 
 def validate_traj_mixer_config(config: Mapping[str, object]) -> None:
@@ -188,7 +188,6 @@ class DeepHealth(nn.Module):
         dist_mode: str = "exponential",  # "exponential", "weibull" or "mixed"
         extra_pool_reduce: str = "mean",
         dropout: float = 0.0,
-        hidden_group: int = 20,
     ):
         super().__init__()
         if target_mode not in ["next_token", "all_future"]:
@@ -243,7 +242,6 @@ class DeepHealth(nn.Module):
                     use_time_rope=False,
                     use_rbf_bias=False,
                     mlp_dropout=dropout,
-                    hidden_group=hidden_group,
                 ) for _ in range(n_hist_layer)
             ])
             self.rope = None
@@ -257,7 +255,6 @@ class DeepHealth(nn.Module):
                     use_time_rope=True,
                     use_rbf_bias=True,
                     mlp_dropout=dropout,
-                    hidden_group=hidden_group,
                 ) for _ in range(n_hist_layer)
             ])
             self.rope = TimeRoPE(n_embd // n_head)