VOGN for regression models

[cyruoff]

Hello,
thank you for that nice repository!
One suggestion:
I would like to try out VOGN for a neural network for regression and wonder if you could put the calculation of probabilities from logits outside of the VOGN optimizer code and make the VOGN optimizer compatible with standard torch optimizers so that it takes the model parameters instead:
optimizer = torchsso.optim.VOGN(model.parameters(), dataset_size=len(train_loader.dataset))
Then it would be compatible with regression models as well.