[Replicate] Simplify code a bit (#98889)
Simplifies the code, such as making self.modules not a list and only a
single module.
Differential Revision: [D44899281](https://our.internmc.facebook.com/intern/diff/D44899281/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/98889
Approved by: https://github.com/mrshenli, https://github.com/yhcharles
