I managed to fix the error. The masking seems to be the culprit.
The line bellow still causes the above mentioned error at a random point during the training process.
loss += torch.sum(criterion(decoder_output, y[:, idx])[valid])/torch.sum(valid)

This line does not cause the aforementioned problem
loss += torch.sum(criterion(decoder_output, y[:, idx])*valid)/torch.sum(valid)

(Here the reduction on the criterion has been set to ‘none’)