diff options
| -rw-r--r-- | text_recognizer/networks/encoders/efficientnet/efficientnet.py | 1 |
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: |