From 55d905e6ee0bb3dbb407f89a626b7cd8c3da5578 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Fri, 29 May 2026 04:38:24 +0300 Subject: [PATCH] Fix gradient: use torch.empty + copy_ to escape inference tensor flag .clone() on an inference tensor inherits the inference flag even inside inference_mode(False). torch.empty() inside that context creates a true normal tensor; copy_() fills it from the inference tensor without flag inheritance. Co-Authored-By: Claude Sonnet 4.6 --- value_function.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/value_function.py b/value_function.py index f53c92cc..8609f230 100644 --- a/value_function.py +++ b/value_function.py @@ -74,11 +74,16 @@ class OnlineValueFunction(nn.Module): def gradient(self, conditioning): """∂reward/∂conditioning — same shape as input.""" mlp_device = next(self.parameters()).device + orig_device = conditioning.device with torch.inference_mode(False), torch.enable_grad(): - c_in = conditioning.detach().clone().float().to(mlp_device).requires_grad_(True) - reward = self.forward(self.compress(c_in).unsqueeze(0)) + # new_empty inside inference_mode(False) creates a normal (non-inference) tensor. + # copy_ fills it from the inference tensor without inheriting the flag. + c_fresh = torch.empty(conditioning.shape, dtype=torch.float32, device=mlp_device) + c_fresh.copy_(conditioning) + c_fresh.requires_grad_(True) + reward = self.forward(self.compress(c_fresh).unsqueeze(0)) reward.backward() - return c_in.grad.to(conditioning.device, conditioning.dtype) + return c_fresh.grad.to(orig_device, conditioning.dtype) def is_ready(self): return len(self.buffer_c) >= self.MIN_SAMPLES