Add MSE minimization along with regularization as an objective

[kanodiaayush]

Also allow for Y to be continuous in [0, 1]; in this case the optimization would be through loss minimization and not through variational inference. Note that we still want the sigmoid to be applied as the final head.

@TianyuDu

[TianyuDu]

Can you have a look to this change: e6bef13
Does it work as you expected?
P.S., this is a preliminary solution to illustrate the idea of potential modifications, I haven't tested it yet.

[kanodiaayush]

Oh. This looks great. This also means our forward function is fully modular, and we can use this infrastructure to compare across models. I will test this out. I will also add l2 regularization in the objective. That should be fine, right?

[TianyuDu]

Great, as long as you only modify bemb_flex_lightning_mse.py and the mse_loss_on_binary_label() function, it will be fine : ).