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 <noreply@anthropic.com>
This commit is contained in:
2026-07-06 13:29:13 +08:00
parent adbbb9d620
commit 8b574bc551
2 changed files with 2 additions and 2 deletions
+1 -1
View File
@@ -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]: