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