fix growing mem usage crnn (#1357)

mindee · Oct 20, 2023 · bc2d3c5 · bc2d3c5
1 parent 4c47ac9
commit bc2d3c5
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/doctr/models/recognition/crnn/pytorch.py b/doctr/models/recognition/crnn/pytorch.py
@@ -128,12 +128,9 @@ def __init__(
         self.feat_extractor = feature_extractor
 
         # Resolve the input_size of the LSTM
-        self.feat_extractor.eval()
-        with torch.no_grad():
+        with torch.inference_mode():
             out_shape = self.feat_extractor(torch.zeros((1, *input_shape))).shape
         lstm_in = out_shape[1] * out_shape[2]
-        # Switch back to original mode
-        self.feat_extractor.train()
 
         self.decoder = nn.LSTM(
             input_size=lstm_in,