Upload app.py

app.py

@@ -87,6 +87,7 @@ class BigramLanguageModel(nn.Module):
             logits = logits[:, -1, :]  # Get the logits for the last token
             probs = nn.functional.softmax(logits, dim=-1)
             idx_next = torch.multinomial(probs, num_samples=1)
+            idx_next = torch.clamp(idx_next, max=60)  # Cap the generated index to max 60
             idx = torch.cat((idx, idx_next), dim=1)
         return idx