diff --git a/bgfilter/segmentation.py b/bgfilter/segmentation.py
index e6269bf..6c6f8b0 100644
--- a/bgfilter/segmentation.py
+++ b/bgfilter/segmentation.py
@@ -62,7 +62,7 @@ class BiRefNetSegmenter:
         """Return a soft foreground mask, float32 0..1, at the input resolution."""
         image = Image.fromarray(rgb.astype(np.uint8), mode="RGB")
         tensor = self.transform(image).unsqueeze(0).to(self.device)
-        with self.torch.no_grad():
+        with self.torch.inference_mode():
             pred = self.model(tensor)[-1].sigmoid()
         pred = pred[0, 0].detach().float().cpu().numpy()
         if pred.shape != rgb.shape[:2]:
diff --git a/bgfilter/vitmatte_infer.py b/bgfilter/vitmatte_infer.py
index d33a265..a12ee95 100644
--- a/bgfilter/vitmatte_infer.py
+++ b/bgfilter/vitmatte_infer.py
@@ -36,7 +36,7 @@ class ViTMatteRunner:
         trimap_image = Image.fromarray(trimap.astype(np.uint8), mode="L")
         inputs = self.processor(images=image, trimaps=trimap_image, return_tensors="pt")
         inputs = {key: value.to(self.device) for key, value in inputs.items()}
-        with self.torch.no_grad():
+        with self.torch.inference_mode():
             outputs = self.model(**inputs)
 
         alpha = outputs.alphas[0, 0].detach().float().cpu().numpy()