Update multi_headed_attn.py

Tencent · Oct 25, 2023 · 3ab1161 · 3ab1161
1 parent ecf369c
commit 3ab1161
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tencentpretrain/layers/multi_headed_attn.py b/tencentpretrain/layers/multi_headed_attn.py
@@ -1,6 +1,7 @@
 import math
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 from tencentpretrain.utils.rope import apply_rotary_emb
 from tencentpretrain.utils.lora import LoraLinear