[Code scan] Handle no-padding pretraining batches in coordinate prediction

[njzjz]

This issue is a result of a Codex global repository scan.

Summary

The pretraining model sets padding_mask to None when a batch has no padding. The coordinate prediction branch later unconditionally calls padding_mask.float(), so batch size 1 or equal-length batches can crash when masked_coord_loss is enabled.

Code references

unimol_tools/unimol_tools/pretrain/unimol.py

Lines 71 to 73 in 4596596

	padding_mask = src_tokens.eq(self.padding_idx)
	if not padding_mask.any():
	padding_mask = None

unimol_tools/unimol_tools/pretrain/unimol.py

Lines 100 to 116 in 4596596

	if self.config.masked_coord_loss > 0:
	coords_emb = src_coord
	if padding_mask is not None:
	atom_num = torch.sum(1 - padding_mask.type_as(x), dim=1).view(
	-1, 1, 1, 1
	) # consider BOS and EOS as part of the object
	else:
	atom_num = src_coord.shape[1]
	delta_pos = coords_emb.unsqueeze(1) - coords_emb.unsqueeze(2)
	attn_probs = self.pair2coord_proj(delta_encoder_pair_rep)
	coord_update = delta_pos / atom_num * attn_probs
	# Mask padding
	pair_coords_mask = (1 - padding_mask.float()).unsqueeze(-1) * (1 - padding_mask.float()).unsqueeze(1)
	coord_update = coord_update * pair_coords_mask.unsqueeze(-1)
	#
	coord_update = torch.sum(coord_update, dim=2)
	encoder_coord = coords_emb + coord_update

Impact

Valid pretraining batches with no padded tokens can raise AttributeError before loss computation. This is especially likely for small datasets, homogeneous molecule sizes, or batch_size=1.

Suggested fix

Keep a separate token-valid mask before nulling padding_mask for the encoder, or branch in the coordinate head and use an all-ones pair mask when padding_mask is None.