From ab462cb07d88cbdfe8837bc2c2f9bb9efec9d2c2 Mon Sep 17 00:00:00 2001
From: BBing <67486385@qq.com>
Date: Tue, 15 Oct 2024 21:54:12 +0800
Subject: [PATCH] Update attn.py

---
 predict/models/attn.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/predict/models/attn.py b/predict/models/attn.py
index cf4785c..d0716aa 100644
--- a/predict/models/attn.py
+++ b/predict/models/attn.py
@@ -31,4 +31,6 @@ class FullAttention(nn.Cell):
         if self.output_attention:
             return V, A
         else:
-            return V, None
\ No newline at end of file
+            return V, None
+
+class ProbAttention(nn.Cell):
\ No newline at end of file