From 3f6be519d50167a04fc206cdbe63b59d7486fd1a Mon Sep 17 00:00:00 2001
From: Sunny Liu <sunny19981005@outlook.com>
Date: Tue, 4 Feb 2025 21:25:13 -0500
Subject: [PATCH] stack

---
 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 1bb8f3d34..3b35f29dd 100644
--- a/src/axolotl/monkeypatch/utils.py
+++ b/src/axolotl/monkeypatch/utils.py
@@ -137,7 +137,7 @@ def get_packed_mask_from_pos_ids(position_ids):
 
         results.append(doc_mask)
 
-    return results
+    return torch.stack(results)
 
 
 def get_seqlens_from_pos_ids(position_ids):