Update on "Renamed bsz to bs for consistency; removed dead code"

some minor cleanups [ghstack-poisoned]
pytorch · May 3, 2024 · 6a3e7b9 · 6a3e7b9
1 parent ccd84c8
commit 6a3e7b9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchtitan/models/llama/model.py b/torchtitan/models/llama/model.py
@@ -420,7 +420,7 @@ def forward(self, tokens: torch.Tensor):
             torch.Tensor: Output logits after applying the Transformer model.
 
         """
-        seqlen = tokens.shape[1]
+        _bs, seqlen = tokens.shape
         h = self.tok_embeddings(tokens)
         self.freqs_cis = self.freqs_cis.to(h.device)
         freqs_cis = self.freqs_cis[0:seqlen]