⛰️[EPIC]: CorrDiff Support & Enhancements

[mnabian]

### Code structuring
- [ ] Split generate.py, train.py and training_loop.py into functions/classes that handle more well-defined tasks
- [ ] Reduce redundancy between the above files
- [x] Clean up handling of config files (https://github.com/NVIDIA/modulus/pull/355)
- [x] Specify clearly in an interface what interface data loaders need to follow, to facilitate adding new loaders (https://github.com/NVIDIA/modulus/pull/376)
- [ ] Same as above, but for models (e.g. U-Nets)
- [ ] Interoperability with diagnostic model code base (TBA)
- [x] Dump checkpoints, logs, and code in separate random directories per run (https://github.com/NVIDIA/modulus/pull/374)

### Optimizations
- [ ] https://github.com/NVIDIA/modulus/issues/424
- [ ] Support distributing generation over time steps as well as ensemble members

### Features
- [ ] Enable per-channel loss
- [ ] https://github.com/NVIDIA/modulus/issues/394
- [ ] https://github.com/NVIDIA/modulus/issues/416
- [ ] https://github.com/NVIDIA/modulus/issues/392
- [ ] Compute loss in inference
- [ ] https://github.com/NVIDIA/modulus/issues/425
- [ ] https://github.com/NVIDIA/modulus/issues/440

### QA
- [x] Make a subset of the CWA dataset

### Misc
- [ ] https://github.com/NVIDIA/modulus/issues/367
- [ ] https://github.com/NVIDIA/modulus/issues/393

### Go-to-market readiness
- [ ] https://github.com/NVIDIA/modulus/issues/407
- [ ] https://github.com/NVIDIA/modulus/issues/408
- [ ] https://github.com/NVIDIA/modulus/issues/409

### Improvements
- [ ] https://github.com/NVIDIA/modulus/issues/414
- [ ] https://github.com/NVIDIA/modulus/issues/415
- [ ] https://github.com/NVIDIA/modulus/issues/417
- [ ] https://github.com/NVIDIA/modulus/issues/441

### Bugs
- [ ] https://github.com/NVIDIA/modulus/issues/418

### Issues
- [ ] https://github.com/NVIDIA/modulus/issues/437

[jleinonen]

#376 is merged so we can probably check off the box "Specify clearly in an interface what interface data loaders need to follow, to facilitate adding new loaders".

[mnabian]

The remaining items will be tracked in this Epic: https://github.com/NVIDIA/modulus/issues/589