Let submodules be collected as args/kwargs (#57840)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/57840
Test Plan: Imported from OSS
Reviewed By: jamesr66a
Differential Revision: D28294984
Pulled By: ansley
fbshipit-source-id: d64fe109a349516da69d2d17f58e42f98af564fd
