From 8b574bc551c55ec00c2db4c93009da7be47d76f9 Mon Sep 17 00:00:00 2001 From: lingh Date: Mon, 6 Jul 2026 13:29:13 +0800 Subject: [PATCH] Switch model inference from no_grad to inference_mode Both call sites convert outputs to numpy immediately, so the stricter inference-mode tensors are safe; saves autograd view/version tracking overhead. Verified end-to-end on CUDA (alpha finite, no NaN). Co-Authored-By: Claude Fable 5 --- bgfilter/segmentation.py | 2 +- bgfilter/vitmatte_infer.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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()