When semantic_embedding_distil is enabled, logits_sem_embed_grid is produced only in the GRID branch and stays as an unpadded JaggedTensor, but downstream SemLoss indexes bb_sem_embed_logits[b][~padding[b]] (expects a padded [B, N, D] tensor like bb_sem_logits). Consider (1) generating the sem-embed logits in both XYZ and GRID modes, and (2) padding/batching them alongside logits via pad_batch(..., additional_feats=...) so the output format matches bb_sem_logits.
Originally posted by @Copilot in #57 (comment)
When
semantic_embedding_distilis enabled,logits_sem_embed_gridis produced only in the GRID branch and stays as an unpadded JaggedTensor, but downstreamSemLossindexesbb_sem_embed_logits[b][~padding[b]](expects a padded[B, N, D]tensor likebb_sem_logits). Consider (1) generating the sem-embed logits in both XYZ and GRID modes, and (2) padding/batching them alongsidelogitsviapad_batch(..., additional_feats=...)so the output format matchesbb_sem_logits.Originally posted by @Copilot in #57 (comment)