Update model.py

model.py

@@ -292,6 +292,10 @@ class GPT(nn.Module):
                 logits[logits < v[:, [-1]]] = -float('Inf')
             probs = F.softmax(logits, dim=-1)
             idx_next = torch.multinomial(probs, num_samples=1)
+            
+            if idx_next.item() == 0: # stop token
+                break
+            
             idx = torch.cat((idx, idx_next), dim=1)
 
         return idx