hllm model

Summary: the HLLM model stacks two LLaMA to perform next item prediction from user's engagement history. Differential Revision: D64089998 Privacy Context Container: L1268898
pytorch · Oct 10, 2024 · 0c3c192 · 0c3c192
1 parent 5de5001
commit 0c3c192
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchtune/modules/transformer.py b/torchtune/modules/transformer.py
@@ -578,7 +578,7 @@ def forward(
             - m_s: max seq len
         """
         # input tensor of shape [b, s]
-        bsz, seq_len = tokens.shape
+        seq_len = tokens.shape[1]
 
         self._validate_inputs(
             seq_len,