From c9ae4a9b461710746b5e62d112edb1037571646c Mon Sep 17 00:00:00 2001 From: felix Date: Thu, 19 Oct 2023 16:19:20 +0200 Subject: [PATCH] fix growing mem usage crnn --- doctr/models/recognition/crnn/pytorch.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/doctr/models/recognition/crnn/pytorch.py b/doctr/models/recognition/crnn/pytorch.py index daf0e56e58..51563a2ade 100644 --- 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,