Support for DataParallel?

[zimmerrol]

Hi, I noticed that the current version does not support parallel models (at least those created using torch.nn.DataParallel) since the forward hook does not differentiate between the different copies of the model and a model wrapped with Inspect will just return the intermediate features of the last copy of the parallelized model to run.

Are you planning on fixing this issue/supporting this use case?

[flavioschneider]

Hi @zimmerrol! Currently, I don't have access to multiple GPUs to try this feature out, nor do I have experience with DataParallel. If your model has a static execution graph, you might be able to output intermediate values using the Extract class before wrapping your model with DataParallel.