fix: dropout torch

mindee · Jul 5, 2021 · b5e3b3e · b5e3b3e
1 parent 32ddedd
commit b5e3b3e
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/doctr/models/recognition/master/tensorflow.py b/doctr/models/recognition/master/tensorflow.py
@@ -398,8 +398,8 @@ def _master(arch: str, pretrained: bool, input_shape: Tuple[int, int, int] = Non
     # Build the model
     model = MASTER(cfg=_cfg, **kwargs)
     # Load pretrained parameters
-    # if pretrained:
-    #     load_pretrained_params(model, default_cfgs[arch]['url'])
+    if pretrained:
+        load_pretrained_params(model, default_cfgs[arch]['url'])
 
     return model
 

diff --git a/doctr/models/recognition/transformer/pytorch.py b/doctr/models/recognition/transformer/pytorch.py
@@ -60,6 +60,8 @@ def __init__(
             ) for _ in range(num_layers)
         ]
 
+        self.dropout = nn.Dropout(dropout)
+
     def forward(
         self,
         x: torch.Tensor,
@@ -73,6 +75,7 @@ def forward(
         x = self.embedding(x)  # (batch_size, target_seq_len, d_model)
         x *= math.sqrt(self.d_model)
         x += self.pos_encoding[:, :seq_len, :]
+        x = self.dropout(x)
 
         # Batch first = False in decoder
         x = x.permute(1, 0, 2)