The only thing special about the forward method is that you can invoke it by calling the model (i.e. m(*inputs) will call m.forward(*inputs)) and that you get hook processing etc.
As such, there is no reason why you would not be able to provide other entry points - I often have a separate predict method for sequence-generating models or sample for density estimation modules.
Whether it is a good idea or not likely depends on your view how closely they are coupled. I’m not entirely sure that it is the case here. On the other hand frameworks like fastai like to stick anything and everything (models, loss functions, learning rate schedule, optimizer, callbacks…) related into one class - but not a nn.Module one.
Best regards
Thomas
Best regards
Thomas