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]:
+1 -1
View File
@@ -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()