GPU OOM error message & auto restart

[thomasckng]

For flowMC, when GPU memory is not enough, a smaller batch size can be chosen to fix it. This can be done without user intervention. Similar treatment can probably be done for other samplers which are currently under development in #105.

[thomasckng]

Currently only the OOM error from JAX is raised without any extra info from Jim or other samplers. Before having auto restart, more informative error is needed, especially for flowMC with parallel tempering, as the error message does not directly tell us it is a OOM issue.