summaryrefslogtreecommitdiff
path: root/text_recognizer/network/convnext/residual.py
diff options
context:
space:
mode:
authorGustaf Rydholm <gustaf.rydholm@gmail.com>2023-08-25 23:19:14 +0200
committerGustaf Rydholm <gustaf.rydholm@gmail.com>2023-08-25 23:19:14 +0200
commit49ca6ade1a19f7f9c702171537fe4be0dfcda66d (patch)
tree20062ed1910758481f3d5fff11159706c7b990c6 /text_recognizer/network/convnext/residual.py
parent0421daf6bd97596703f426ba61c401599b538eeb (diff)
Rename and add flash atten
Diffstat (limited to 'text_recognizer/network/convnext/residual.py')
-rw-r--r--text_recognizer/network/convnext/residual.py16
1 files changed, 16 insertions, 0 deletions
diff --git a/text_recognizer/network/convnext/residual.py b/text_recognizer/network/convnext/residual.py
new file mode 100644
index 0000000..dfc2847
--- /dev/null
+++ b/text_recognizer/network/convnext/residual.py
@@ -0,0 +1,16 @@
+"""Generic residual layer."""
+from typing import Callable
+
+from torch import Tensor, nn
+
+
+class Residual(nn.Module):
+ """Residual layer."""
+
+ def __init__(self, fn: Callable) -> None:
+ super().__init__()
+ self.fn = fn
+
+ def forward(self, x: Tensor) -> Tensor:
+ """Applies residual fn."""
+ return self.fn(x) + x