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 = contraint

  def forward(self, inputs):
    predictions = self.model(inputs)
    loss_1, loss_2 = calculate_loss(predictions)
    total_loss = loss_1 + constraint(loss_2)
    return total_loss

ddp = DistributedDataParallel(
  WrapperModule(model, constraints), 
  device_ids=[device_rank], 
  find_unused_parameters=True
)
model_optimiser = optimiser(model.parameters(), ...)
constraint_optimiser = optimiser(constraint.parameters(), ...)

ddp(inputs).backward()
model_optimiser.step()
constraint_optimiser.step()

BTW, I noticed that the code uses two different processes groups for model and constraints. When using NCCL backend, this could lead to deadlock. As of 2.8, NCCL requires using only one communicator for each CUDA device at any given time instance (i.e., not thread-safe)