daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix Downsample3D spatial reflect pad to (W,H) dims only

Browse source
This commit is contained in:
Daniel Maddern 2026-08-19 21:41:16 +07:00
parent 136297f4ae
commit b557cba173
1 changed files with 3 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
src/h3_blackwell_runtime/vae_encoder.py
View file

@ -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)