diff --git a/src/h3_blackwell_runtime/vae_encoder.py b/src/h3_blackwell_runtime/vae_encoder.py index 29d0d64..869fa81 100644 --- a/src/h3_blackwell_runtime/vae_encoder.py +++ b/src/h3_blackwell_runtime/vae_encoder.py @@ -93,7 +93,9 @@ def _resnet(x, p): def _downsample(x, p): if p["space"] == 2: - x = F.pad(x, (0, 1, 0, 1, 0, 0), mode="reflect") + # Reference Downsample3D pads H and W by +1 reflect before the conv. + # F.pad for 5D (B,C,T,H,W) uses (W_l, W_r, H_l, H_r, T_l, T_r). + x = F.pad(x, (1, 1, 1, 1, 0, 0), mode="reflect") return _causal_conv3d(x, p["w"], p["b"], kernel_size=3, stride=(p["time"], p["space"], p["space"]), spatial_padding=0)