fix bug when generate

modeling_baichuan.py

@@ -221,10 +221,10 @@ class Attention(nn.Module):
         #                  head_dim=self.head_dim, n_head=self.num_heads)
         # q, k, v = proj
         if past_key_value is None:
+            self.last_k = k[:,-1:]   
+            self.last_v = v[:,-1:]
             k = custom_convolution(k, self.K)
             v = custom_convolution(v, self.V)
-            self.last_k = k[:,-1:]
-            self.last_v = v[:,-1:]
         else:
             self.last_k,k = k, self.K[:,:1]*self.last_k + self.K[:,1:]*k
             self.last_v,v = v, self.V[:,:1]*self.last_v + self.V[:,1:]*v