fix reshape in smooth lddt loss

[grfrederic]

I believe the reshape on line 161 in smooth_lddt_loss is incorrect and will misalign mask and eps if both batch and multiplicity are larger than 1.

The reason is that if some xs has shape (batch, d1, d2, ...) then repeat_interleave has the following property:

ys = torch.repeat_interleave(multiplicity, 0)  # shape (batch * multiplicity, d1, d2, ...)
xs == ys.reshape(batch, multiplicity, d1, d2, ...)

That is, since repeat_interleave repeats every element of the batch a given number of times, not the whole batch, reshape should put it on dim=1. However, on line 161 in the current code, the reshape puts multiplicity on dim=0, and averages over it.

Additionally, since the resulting averaged eps (shape (batch, d1, d2, ...)) doesn't match the mask (which has shape (batch * multiplicity, d1, d2, ...)), eps is interleaved again on line 166. If lines 161-162 were to be changed to .view(B // multiplicity, multiplicity, N, N).mean(dim=1), the repeat_interleave on line 166 simply undoes this reshape and average. Thus, I suggest simply removing these lines.