[PyTorch] replace no_grad with inference_mode

doctr/models/classification/predictor/pytorch.py

@@ -33,7 +33,7 @@ def __init__(
         self.pre_processor = pre_processor
         self.model = model.eval()
 
-    @torch.no_grad()
+    @torch.inference_mode()
     def forward(
         self,
         crops: List[Union[np.ndarray, torch.Tensor]],

doctr/models/detection/predictor/pytorch.py

@@ -32,7 +32,7 @@ def __init__(
         self.pre_processor = pre_processor
         self.model = model.eval()
 
-    @torch.no_grad()
+    @torch.inference_mode()
     def forward(
         self,
         pages: List[Union[np.ndarray, torch.Tensor]],

doctr/models/kie_predictor/pytorch.py

@@ -59,7 +59,7 @@ def __init__(
         self.detect_orientation = detect_orientation
         self.detect_language = detect_language
 
-    @torch.no_grad()
+    @torch.inference_mode()
     def forward(
         self,
         pages: List[Union[np.ndarray, torch.Tensor]],

doctr/models/predictor/pytorch.py

@@ -59,7 +59,7 @@ def __init__(
         self.detect_orientation = detect_orientation
         self.detect_language = detect_language
 
-    @torch.no_grad()
+    @torch.inference_mode()
     def forward(
         self,
         pages: List[Union[np.ndarray, torch.Tensor]],

doctr/models/recognition/predictor/pytorch.py

@@ -40,7 +40,7 @@ def __init__(
         self.dil_factor = 1.4  # Dilation factor to overlap the crops
         self.target_ar = 6  # Target aspect ratio
 
-    @torch.no_grad()
+    @torch.inference_mode()
     def forward(
         self,
         crops: Sequence[Union[np.ndarray, torch.Tensor]],

references/classification/latency_pytorch.py

@@ -19,7 +19,7 @@
 from doctr.models import classification
 
 
-@torch.no_grad()
+@torch.inference_mode()
 def main(args):
     device = torch.device("cuda:0" if args.gpu else "cpu")
 

references/detection/evaluate_pytorch.py

@@ -26,7 +26,7 @@
 from doctr.utils.metrics import LocalizationConfusion
 
 
-@torch.no_grad()
+@torch.inference_mode()
 def evaluate(model, val_loader, batch_transforms, val_metric, amp=False):
     # Model in eval mode
     model.eval()

references/detection/latency_pytorch.py

@@ -19,7 +19,7 @@
 from doctr.models import detection
 
 
-@torch.no_grad()
+@torch.inference_mode()
 def main(args):
     device = torch.device("cuda:0" if args.gpu else "cpu")
 

references/obj_detection/latency_pytorch.py

@@ -19,7 +19,7 @@
 from doctr.models import obj_detection
 
 
-@torch.no_grad()
+@torch.inference_mode()
 def main(args):
     device = torch.device("cuda:0" if args.gpu else "cpu")
 

references/recognition/evaluate_pytorch.py

@@ -22,7 +22,7 @@
 from doctr.utils.metrics import TextMatch
 
 
-@torch.no_grad()
+@torch.inference_mode()
 def evaluate(model, val_loader, batch_transforms, val_metric, amp=False):
     # Model in eval mode
     model.eval()

references/recognition/latency_pytorch.py

@@ -19,7 +19,7 @@
 from doctr.models import recognition
 
 
-@torch.no_grad()
+@torch.inference_mode()
 def main(args):
     device = torch.device("cuda:0" if args.gpu else "cpu")
 

scripts/detect_artefacts.py

@@ -37,7 +37,7 @@ def plot_predictions(image, boxes, labels):
     plt.show()
 
 
-@torch.no_grad()
+@torch.inference_mode()
 def main(args):
     print(args)