Extended support for wrapping container modules (ModuleList, ModuleDict, Sequential) in Fabric

Description & Motivation

Fabric's setup methods wrap the user's nn.Module in a FabricModule. This one forwards all calls to the original module with a few exceptions where we add additional logic before and after.

There are three additional special module types in PyTorch: Sequential, ModuleList, and ModuleDict. These have additional methods like __getitem__ for advanced indexing. The FabricModule does not support handling these:

module = torch.nn.Sequential(torch.nn.Linear(1, 1), torch.nn.ReLU())
fabric_module = fabric.setup(module)
module[0]  # works
fabric_module[0]  # fails

module = torch.nn.ModuleDict({"module_key": torch.nn.Linear(1, 1)})
fabric_module = fabric.setup(module)
module["module_key"]  # works
fabric_module["module_key"]  # fails

Pitch

Implement the dunder methods (or a subset of them) for the FabricModule wrapper.

Alternatives

Keep as is, but document that the desired functionality can be accessed via fabric_module.module.

Additional context

No response

cc @borda

Start by locating Fabric's setup methods and the FabricModule wrapper, then compare its behavior with PyTorch's Sequential, ModuleList, and ModuleDict container methods linked in the issue. Decide which dunder methods are needed for indexing and verify that wrapped containers support the examples shown in the issue.