You may also want to see this thread. When using the PhysicsInformedNN interface, you could use different NN to represent each solution. When using the NNODE interface, you can use a single NN with multiple outputs to represent each solution. I see that you are not using either interface and rolling your own solution, so I am not sure if my reply really helps.