From 1f459ba19422593de325983040e176f97cf4ffc0 Mon Sep 17 00:00:00 2001 From: aktersnurra Date: Thu, 20 Aug 2020 22:18:35 +0200 Subject: A lot of stuff working :D. ResNet implemented! --- src/text_recognizer/models/character_model.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) (limited to 'src/text_recognizer/models/character_model.py') diff --git a/src/text_recognizer/models/character_model.py b/src/text_recognizer/models/character_model.py index 0a0ab2d..0fd7afd 100644 --- a/src/text_recognizer/models/character_model.py +++ b/src/text_recognizer/models/character_model.py @@ -44,6 +44,7 @@ class CharacterModel(Model): self.tensor_transform = ToTensor() self.softmax = nn.Softmax(dim=0) + @torch.no_grad() def predict_on_image( self, image: Union[np.ndarray, torch.Tensor] ) -> Tuple[str, float]: @@ -64,10 +65,9 @@ class CharacterModel(Model): # If the image is an unscaled tensor. image = image.type("torch.FloatTensor") / 255 - with torch.no_grad(): - # Put the image tensor on the device the model weights are on. - image = image.to(self.device) - logits = self.network(image) + # Put the image tensor on the device the model weights are on. + image = image.to(self.device) + logits = self.network(image) prediction = self.softmax(logits.data.squeeze()) -- cgit v1.2.3-70-g09d2