From 85752cdfc955550aaa6d045000c96723fca0420e Mon Sep 17 00:00:00 2001 From: Sunny Liu Date: Wed, 22 Jan 2025 21:33:38 -0500 Subject: [PATCH] mask expansion --- src/axolotl/monkeypatch/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/axolotl/monkeypatch/utils.py b/src/axolotl/monkeypatch/utils.py index da85aa9cc..c360cf94e 100644 --- a/src/axolotl/monkeypatch/utils.py +++ b/src/axolotl/monkeypatch/utils.py @@ -229,7 +229,7 @@ def mask_2d_to_4d( tgt_len = tgt_len if tgt_len is not None else src_len # mask = mask.unsqueeze(1).unsqueeze(2) - mask = mask[:, None, :].expand(bsz, 1, tgt_len, src_len) + mask = mask[None, None, :, :].expand(bsz, 1, tgt_len, src_len) # Create a binary mask from the original mask where zeros remain zeros and all other values are set to one binary_mask = torch.where(