📚[DOC]: improve the UNet example in the documentation of Module #1178
Description
How would you describe the priority of this documentation request
Critical (currently preventing usage)
Is this for new documentation, or an update to existing docs?
New
Describe the incorrect/future/missing documentation
As noted by @nbren12 the UNet example in the documentation is unclear. This part of the documentation should clearly illustrate the use of physicsnemo.Modules and the unique functionalities they provide (checkpointing, versioning system, backward compatibility, etc...). Instead the example is cluttered with details of layers implementation that are irrelevant to PhysicsNeMo's Module class. The example also shows examples of training loop and do so on, which is not directly related to defining and using the Module class and results in blurring the message that the documentation should convey.
In addition, the documentation should detail the mechanics used to register new Module's subclasses in the registry, as the relation between the name attribute in the Module's MetaData and the class.__name__ is not clear. In particular, what name is used to lookup the model in the registry is not clear and should be better explained.