From 0c3c192fbaed7f422d95ea2fb6e28417e5ef4f76 Mon Sep 17 00:00:00 2001
From: Qiang Zhang <qiangzhang@meta.com>
Date: Thu, 10 Oct 2024 12:16:58 -0700
Subject: [PATCH] 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
---
 torchtune/modules/transformer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torchtune/modules/transformer.py b/torchtune/modules/transformer.py
index 16a66f5aa6..440cba3001 100644
--- 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,