The pseudo code is something like this
output = model(input)
output2 = model(input2)
final_output = model.module.function(output, output2)
loss = loss_function(final_output)
optimizer.zero_grad()
loss.backward()
optimizer.step()

Would this be fine? The custom function is just a MLP to classify something. It does not change anything, but I want it to get updated when I call my optimizer.step()