bigyeet:

Is this right, or do I have to write a custom DataParallel wrapper that has scatter, gather, etc methods? If so, how would I do it?


It depends on what you expected reset_hidden_state to achieve. Below is what happens in EVERY forward pass when you use DataParallel.

split input data
replicate model to all devices
feed input data splits to all model replicas
gather outputs from all replicas
done with forward

After the forward pass, the autograd graph actually contains multiple model replicas. It looks sth like
original model <- scatter <- model replicas <- replica output <- gather <- final output.
So in your above use case, if reset_hidden_state has any side effect that you would like to apply to the backward pass, it will only apply to the original model, not to model replicas. But if you are only trying to clear some states for the next forward pass, it should work.