[Flash Attention] Fix flash attention integration

src/transformers/generation/utils.py

@@ -681,12 +681,12 @@ def prepare_inputs_for_generation(
             tensor_kws = {"dtype": torch.int32, "device": self.device}
             pos = model_inputs["position_ids"][:, -1]
 
-            cu_seq_lens_k = torch.cat([torch.zeros(1, **tensor_kws), pos.cumsum(0).add(1)], 0)
+            cu_seq_lens_k = torch.cat([torch.zeros(1, **tensor_kws), pos.cumsum(0).add(1)], dim=0).to(**tensor_kws)
             max_length_k = int(pos.max()) + 1
 
             bs, seq_len = input_ids.size()
             q_len = torch.ones(bs, **tensor_kws) if seq_len == 1 else pos.to(torch.int32).add(1)
-            cu_seq_lens_q = torch.cat([torch.zeros(1, **tensor_kws), q_len.cumsum(0)], 0)
+            cu_seq_lens_q = torch.cat([torch.zeros(1, **tensor_kws), q_len.cumsum(0)], dim=0).to(**tensor_kws)
             max_length_q = int(q_len.max())
 
             model_inputs.update(