Scale TrajMixer hidden width with head count

train_next_step.py

@@ -84,7 +84,6 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--time_mode", type=str, default="relative",
                         choices=["relative", "absolute"])
     parser.add_argument("--dropout", type=float, default=0.0)
-    parser.add_argument("--hidden_group", type=int, default=20)
 
     parser.add_argument("--target_mode", type=str, default="uts",
                         choices=["delphi2m", "uts"])
@@ -166,7 +165,6 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
         time_mode=args.time_mode,
         dist_mode="exponential",
         dropout=args.dropout,
-        hidden_group=args.hidden_group,
     )