summaryrefslogtreecommitdiff
path: root/text_recognizer/networks/encoders/efficientnet/efficientnet.py
diff options
context:
space:
mode:
authorGustaf Rydholm <gustaf.rydholm@gmail.com>2021-11-07 23:48:15 +0100
committerGustaf Rydholm <gustaf.rydholm@gmail.com>2021-11-07 23:48:15 +0100
commit8728fc193af64de1eb70c0b24b1c3e80d97073dc (patch)
tree29a736dfb98b310840c28a52f2433a6ea628b012 /text_recognizer/networks/encoders/efficientnet/efficientnet.py
parentec68cacdcd0a85f0cf5b39529f4a93cfde7db838 (diff)
Add dropout to final layer in efficientnet
Diffstat (limited to 'text_recognizer/networks/encoders/efficientnet/efficientnet.py')
-rw-r--r--text_recognizer/networks/encoders/efficientnet/efficientnet.py1
1 files changed, 1 insertions, 0 deletions
diff --git a/text_recognizer/networks/encoders/efficientnet/efficientnet.py b/text_recognizer/networks/encoders/efficientnet/efficientnet.py
index 7023297..32934d5 100644
--- a/text_recognizer/networks/encoders/efficientnet/efficientnet.py
+++ b/text_recognizer/networks/encoders/efficientnet/efficientnet.py
@@ -101,6 +101,7 @@ class EfficientNet(nn.Module):
momentum=self.bn_momentum,
eps=self.bn_eps,
),
+ nn.Dropout(p=self.arch[-1]),
)
def extract_features(self, x: Tensor) -> Tensor: