Update functional.py

morgan-heisler · web-flow · commit a050b21b4cad · 2026-04-20T15:47:16.000-07:00
update 2D padded format as well
diff --git a/areal/utils/functional/functional.py b/areal/utils/functional/functional.py
@@ -426,7 +426,7 @@ def apply_rejection_sampling(
                     if config.lower is not None:
                         token_oor = token_oor | (token_ratio < config.lower)
                     loss_mask = loss_mask * (~token_oor).to(loss_mask.dtype)
-                    behave_imp_weight = behave_imp_weight * (~token_oor).to(
+                    behave_imp_weight = token_ratio * (~token_oor).to(
                         behave_imp_weight.dtype
                     )
                     

Original file line number	Diff line number	Diff line change
`@@ -426,7 +426,7 @@ def apply_rejection_sampling(`
`426`	`426`	`if config.lower is not None:`
`427`	`427`	`token_oor = token_oor \| (token_ratio < config.lower)`
`428`	`428`	`loss_mask = loss_mask * (~token_oor).to(loss_mask.dtype)`
`429`		`- behave_imp_weight = behave_imp_weight * (~token_oor).to(`
	`429`	`+ behave_imp_weight = token_ratio * (~token_oor).to(`
`430`	`430`	`behave_imp_weight.dtype`
`431`	`431`	`)`
`432`	`432`