replace no_grad with inference_mode (#1323)

mindee · Sep 26, 2023 · 56395ba · 56395ba
1 parent 9fa5e78
commit 56395ba
Show file tree

Hide file tree

Showing 12 changed files with 12 additions and 12 deletions.
diff --git a/doctr/models/classification/predictor/pytorch.py b/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]],

diff --git a/doctr/models/detection/predictor/pytorch.py b/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]],

diff --git a/doctr/models/kie_predictor/pytorch.py b/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]],

diff --git a/doctr/models/predictor/pytorch.py b/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]],

diff --git a/doctr/models/recognition/predictor/pytorch.py b/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]],

diff --git a/references/classification/latency_pytorch.py b/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")
 

diff --git a/references/detection/evaluate_pytorch.py b/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()

diff --git a/references/detection/latency_pytorch.py b/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")
 

diff --git a/references/obj_detection/latency_pytorch.py b/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")
 

diff --git a/references/recognition/evaluate_pytorch.py b/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()

diff --git a/references/recognition/latency_pytorch.py b/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")
 

diff --git a/scripts/detect_artefacts.py b/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)