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:
@@ -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]:
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user