If model is a DistributedDataParallel (DDP) instance, this won’t work. Because setup some internal states at the end of the forward pass, and does not work if you call forward twice without a backward in between.
However, this can be easily solve by wrapping the two forward and the function invocation into a wrapper model, and then pass that wrapper model to DDP, sth like:
class WrapperModel(nn.Module):
  def __init__(self, model) :
    super(WrapperModel, self).__init__()
    self.model = model

  def forward(input, input2): 
    output = model(input)
    output2 = model(input2)
    final_output = model.module.function(output, output2)
    return final_output

ddp = DistributedDataParallel(WrapperModel(model).to(device), device_ids=[device])

final_output = ddp.forward(input, input2)
loss = loss_function(final_output)
optimizer.zero_grad()
loss.backward()
optimizer.step()