Fix cont loss weighting for filtered cont dims

2026-01-28 22:27:23 +08:00
parent 8db286792e
commit 6fb53dd5c1
1 changed files with 1 additions and 1 deletions
--- a/example/train.py
+++ b/example/train.py
@@ -355,7 +355,7 @@ def main():

            if config.get("cont_loss_weighting") == "inv_std":
                weights = torch.tensor(
-                    [1.0 / (float(raw_std[c]) ** 2 + float(config.get("cont_loss_eps", 1e-6))) for c in cont_cols],
+                    [1.0 / (float(raw_std[c]) ** 2 + float(config.get("cont_loss_eps", 1e-6))) for c in model_cont_cols],
                    device=device,
                    dtype=eps_pred.dtype,
                ).view(1, 1, -1)