From 3fc943704f8232d296e1d8f5869ea3d40541b6dc Mon Sep 17 00:00:00 2001 From: Riccardo Rota <144956381+Riccardo-Rota@users.noreply.github.com> Date: Mon, 23 Mar 2026 10:37:08 +0100 Subject: [PATCH] Fix slicing of latent encodings for merging Latent encodings were sliced as x_latent0_encodings[: batch_size * x0_tile_size], changed into x_latent0_encodings[: x0_tile_size] because x0_tile_size already accounts for batch size. For example, if batch_size=10, the original x_latent0_encodings has shape [350, 1024, 24, 24] and x0_tile_size=250. If we multiply again by batch_size=10, we slice up to index 2500, making the slicing uneffective. Same happens for x_latent1_encodings --- src/sharp/models/encoders/spn_encoder.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/sharp/models/encoders/spn_encoder.py b/src/sharp/models/encoders/spn_encoder.py index 324a359..d0e6c07 100644 --- a/src/sharp/models/encoders/spn_encoder.py +++ b/src/sharp/models/encoders/spn_encoder.py @@ -255,7 +255,7 @@ def forward(self, x: torch.Tensor) -> list[torch.Tensor]: patch_intermediate_features[self.patch_intermediate_features_ids[0]] # type:ignore[index] ) x_latent0_features = merge( - x_latent0_encodings[: batch_size * x0_tile_size], + x_latent0_encodings[: x0_tile_size], batch_size=batch_size, padding=padding, ) @@ -264,7 +264,7 @@ def forward(self, x: torch.Tensor) -> list[torch.Tensor]: patch_intermediate_features[self.patch_intermediate_features_ids[1]] # type:ignore[index] ) x_latent1_features = merge( - x_latent1_encodings[: batch_size * x0_tile_size], + x_latent1_encodings[: x0_tile_size], batch_size=batch_size, padding=padding, )