From db5c766a8e6684c82dbde3c82634bf5d98e8f877 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Fri, 29 May 2026 03:57:30 +0300 Subject: [PATCH] Fix value function: wrap entire train_on in inference_mode(False) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Buffer tensors were created inside inference_mode — detach().cpu() doesn't escape the inference flag. Wrapping the whole method ensures the clone and all subsequent tensor ops produce normal autograd-tracked tensors. Co-Authored-By: Claude Sonnet 4.6 --- value_function.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/value_function.py b/value_function.py index 090561bd..606cc60a 100644 --- a/value_function.py +++ b/value_function.py @@ -52,19 +52,19 @@ class OnlineValueFunction(nn.Module): def train_on(self, conditioning, reward): """Add a new (conditioning, reward) sample and do a few training steps.""" - c = self.compress(conditioning).detach().cpu() - self.buffer_c.append(c) - self.buffer_r.append(float(reward)) - if len(self.buffer_c) > self.BUFFER_SIZE: - self.buffer_c = self.buffer_c[-self.BUFFER_SIZE:] - self.buffer_r = self.buffer_r[-self.BUFFER_SIZE:] - if len(self.buffer_c) < 2: - return - device = next(self.parameters()).device with torch.inference_mode(False), torch.enable_grad(): + c = self.compress(conditioning.clone()).detach().cpu() + self.buffer_c.append(c) + self.buffer_r.append(float(reward)) + if len(self.buffer_c) > self.BUFFER_SIZE: + self.buffer_c = self.buffer_c[-self.BUFFER_SIZE:] + self.buffer_r = self.buffer_r[-self.BUFFER_SIZE:] + if len(self.buffer_c) < 2: + return + device = next(self.parameters()).device for _ in range(self.TRAIN_STEPS): idx = random.sample(range(len(self.buffer_c)), min(self.BATCH_SIZE, len(self.buffer_c))) - c_b = torch.stack([self.buffer_c[i] for i in idx]).to(device).clone() + c_b = torch.stack([self.buffer_c[i] for i in idx]).to(device) r_b = torch.tensor([self.buffer_r[i] for i in idx], device=device).unsqueeze(1) self.optimizer.zero_grad() F.mse_loss(self.forward(c_b), r_b).backward()