From 32438179cda0e69a8e6743c6e7c6af91ff458eb0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Oliver=20D=C3=BCrr?= Date: Thu, 2 Jul 2026 18:10:59 +0200 Subject: [PATCH] Fix sample(do=...) to follow the model dtype torch.full defaulted to float32, crashing sampling under do on a float64 model (flow.double(), e.g. for tight cross-implementation parity checks). Follow self._dtype like every other value tensor. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01H1WDjRpYVvCvFG99Lirnsi --- src/tramdag/flow.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/tramdag/flow.py b/src/tramdag/flow.py index 5601aa4..06fe455 100644 --- a/src/tramdag/flow.py +++ b/src/tramdag/flow.py @@ -629,7 +629,8 @@ def sample(self, n: int | None = None, *, do: dict[str, float] | None = None, values: dict[str, Tensor] = {} for name in self.order: if name in do: - values[name] = torch.full((n,), float(do[name]), device=self.device) + values[name] = torch.full((n,), float(do[name]), + dtype=self._dtype, device=self.device) continue node = self.nodes[name] feats = self._features({p: values[p] for p in node.parents})