github
/
InternLM
mirror of https://github.com/InternLM/InternLM
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

merge upstream develop

pull/407/head
yingtongxiong 2023-10-09 21:35:52 +08:00
parent a8dea6313f
commit 1b7935dd98
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
internlm/model/multi_head_attention.py
View File

@ -6,6 +6,7 @@ from typing import Any, Optional, Tuple
import torch import torch
import torch.distributed as dist import torch.distributed as dist
import torch.nn.functional as F
from einops import rearrange from einops import rearrange
from flash_attn.modules.mha import ( from flash_attn.modules.mha import (
CrossAttention, CrossAttention,