This is an example of how it should work:
model_0 = create_model()
summed_losses = 0
for _ in range(epochs):
  models_array = [deepcopy(model) for _ in range N] # get N copies of the main model
  for i in range(N):
    local_optimizer = optimizer(models_array[i].parameters(), lr)
    local_loss = loss(models_array[i](data[i]), labels[i]) 
    #each model is trained on different data

    local_optimizer.zero_grad()
    local_loss.backward()
    local_optimizer.step()

    with torch.no_grad():    
      summed_losses += loss(models_array[i](eval_data[i]), eval_labels[i]) #collect sum of losses for the main model

  main_optimizer = optimizer(model_0.parameters(), lr)
  main_loss = summed_losses #this is the part that I want to figure out. How do I assign summed_losses as the loss for my main model
  
  main_optimizer.zero_grad()
  main_loss.backward() #trick the main model to use the summed loss as its own
  main_optimizer.step() 

so at every iteration, I initiate N copies of the main model, then they train on different data and update their weights independently. Then I evaluate those N models and get a summation of their losses, and use that loss as the main model’s loss.