From 1b7935dd98d7879ae7effd1723ffa70e32869c5e Mon Sep 17 00:00:00 2001
From: yingtongxiong <974106207@qq.com>
Date: Mon, 9 Oct 2023 21:35:52 +0800
Subject: [PATCH] merge upstream develop

---
 internlm/model/multi_head_attention.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/internlm/model/multi_head_attention.py b/internlm/model/multi_head_attention.py
index 287a0e2..49578d7 100644
--- a/internlm/model/multi_head_attention.py
+++ b/internlm/model/multi_head_attention.py
@@ -6,6 +6,7 @@ from typing import Any, Optional, Tuple
 
 import torch
 import torch.distributed as dist
+import torch.nn.functional as F
 from einops import rearrange
 from flash_attn.modules.mha import (
     CrossAttention,