Self_training_fix

[changbiHub]

Fix bugs in EncoderDecoderTrainer and EncoderDecoderLoss

I identified two bugs while using EncoderDecoderTrainer for pretraining:

Issue 1: Negative Loss Values in EncoderDecoderLoss

When calculating the loss, if x_true_stds is zero, it's replaced with the corresponding value from x_true_means. If these mean values are negative, they cause the loss to become negative, which is problematic for optimization.

The problematic code:

x_true_stds = torch.std(x_true, dim=0) ** 2
x_true_stds[x_true_stds == 0] = x_true_means[x_true_stds == 0]

Issue 2: Inconsistent Return Order in _forward_tabnet

The _forward_tabnet method returns values in a different order compared to other encoder-decoder models:

• Other models: x_embed, x_embed_rec, mask
• TabNet: x_embed_rec, x_embed, mask

This inconsistency causes incorrect inputs to the loss function when using TabNet with EncoderDecoderTrainer.

Solution

For Issue 1:

Modified the EncoderDecoderLoss to use the absolute value of means when replacing zero standard deviations:

x_true_stds[x_true_stds == 0] = torch.abs(x_true_means[x_true_stds == 0])

This ensures that the scaling factor remains positive, which is conceptually correct since standard deviations are always non-negative.

For Issue 2:

Changed the return order in _forward_tabnet to match the convention used by other models:

return x_embed, x_embed_rec, mask  # Previously: x_embed_rec, x_embed, mask

[jrzaurin]

thanks for opening the PR. Next time (I hope there is a next time :) ), if you could please open it from a brach in the repo would be great, as it would go through the tests. Now it has broken a few tests. I will try to fix them

Thanks anyway!

[jrzaurin]

ok, errors have to do with pytorch 2.6.

I will open an issue

[changbiHub]

Sorry about that, first time PR.