I used a function outside of the model class to perform the forward pass. It seems DDP requires a forward method nested in the nn.Module, which is not stated in the documentation.
After changing my code to include the forward method the error is the following:
RuntimeError: Expected to mark a variable ready only once. This error is caused by one of the following reasons: 1) Use of a module parameter outside the 'forward' function. Please make sure model parameters are not shared across multiple concurrent forward-backward passes2) Reused parameters in multiple reentrant backward passes. For example, if you use multiple 'checkpoint' functions to wrap the same part of your model, it would result in the same set of parameters been used by different reentrant backward passes multiple times, and hence marking a variable ready multiple times. DDP does not support such use cases yet.
The code works fine on a single GPU. So I do not quite understand the error, since backpropagating twice through the same graph without retain_graph=True would not be possible anyway.