Allow to specify num_mols in OutputModel

[raimis]

Alternative to #214

Just specify the number of molecules to make CUDA Graphs work:

model = load_model(model_file, num_mols=1)

[RaulPPelaez]

I am scared to add another parameter like this. If one saves this model it will lock it to that num_mol. On the other hand this parameter exists solely for the purpose of CUDA graphs, but it can be misleading to an user who might think this parameter has to be set.

If we do this we have to document it really good.

[raimis]

If #214 works, we can drop this.

[raimis]

The solution in #214 is better.