diff --git a/src/unilab/training/tensor_diagnostics.py b/src/unilab/training/tensor_diagnostics.py index 080ebb324..9e19c7579 100644 --- a/src/unilab/training/tensor_diagnostics.py +++ b/src/unilab/training/tensor_diagnostics.py @@ -128,10 +128,18 @@ def capture_due(self, step: int) -> bool: return self._cfg.enabled and interval > 0 and step % interval == 0 def reset_verdict(self) -> None: - if self._device_bad_rows is not None: - self._device_bad_rows.zero_() - if self._device_verdict is not None: - self._device_verdict.zero_() + # The device latches may be inference tensors: rsl_rl wraps the whole + # rollout in torch.inference_mode(), so with deferred validation the + # buffers are typically first allocated inside that context. In-place + # zeroing of an inference tensor is only legal under inference mode, + # and poll sites (cycle boundaries, close) run outside it, so run the + # reset under inference mode explicitly; it is a no-op on the + # autograd bookkeeping either way. + with torch.inference_mode(): + if self._device_bad_rows is not None: + self._device_bad_rows.zero_() + if self._device_verdict is not None: + self._device_verdict.zero_() def _validate_finite_float(self, value: torch.Tensor, label: str) -> None: if not isinstance(value, torch.Tensor): diff --git a/tests/training/test_tensor_diagnostics.py b/tests/training/test_tensor_diagnostics.py index e19fc83b3..87a2995ab 100644 --- a/tests/training/test_tensor_diagnostics.py +++ b/tests/training/test_tensor_diagnostics.py @@ -200,3 +200,20 @@ def test_training_routes_guard_and_shared_finite_boundary() -> None: cfg.training.finite_validation_interval = 0 with pytest.raises(ValueError, match="positive integer"): finite_validation_env_overrides(cfg) + + +def test_reset_verdict_allows_inference_tensor_latches(tmp_path: Path) -> None: + """Latches allocated inside an rsl_rl rollout must survive an outside reset. + + rsl_rl wraps the whole rollout in ``torch.inference_mode()``; with + deferred validation the guard's device buffers are typically first + allocated there, and ``close()`` resets them outside the context. The + reset must not fail on inference tensors. + """ + guard = _guard(tmp_path) + obs = {"obs": torch.zeros((3, 2))} + reward = torch.zeros(3) + with torch.inference_mode(): + assert guard.check(obs, reward, defer=True) is None + assert guard.device_verdict is not None + guard.reset_verdict()