Solved by mrshenli in post #2 
              
              
                Hey @ClaartjeBarkhof, can you try wrapping your model and constraint into one nn.Module first and then wrap that module with DistributedDataParallel. Sth like: 

class WrapperModule(nn.Module):
  def __init__(self, model, constraint):
    super().__init__()
    self.model = model
    self.constraint…