Scale TrajMixer hidden width with head count

train_all_future.py

@@ -90,7 +90,6 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--dist_mode", type=str, default="exponential",
                         choices=["exponential", "weibull", "mixed"])
     parser.add_argument("--dropout", type=float, default=0.0)
-    parser.add_argument("--hidden_group", type=int, default=20)
 
     parser.add_argument("--batch_size", type=int, default=128)
     parser.add_argument("--base_lr", type=float, default=3e-4)
@@ -161,7 +160,6 @@ def build_model(args: argparse.Namespace, dataset: AllFutureHealthDataset) -> De
         time_mode=args.time_mode,
         dist_mode=args.dist_mode,
         dropout=args.dropout,
-        hidden_group=args.hidden_group,
     )