code: def local_aggregation_train(
      model, optimizer, dataloader, isCUDA, mini_batch_size
  ):
      mseloss = nn.MSELoss()
      model.train()
      for data, target in dataloader:
          if isCUDA:
              data, target = data.cuda(), target.cuda()
          optimizer.zero_grad()
          for i in range(0, len(data), mini_batch_size):
              data_batch = data[i : i + mini_batch_size]
              target_batch = target[i : i + mini_batch_size]
              y = model(data_batch)
              loss = mseloss(y, target_batch)
              loss.div_(len(data) / mini_batch_size)
              loss.backward()
          optimizer.step()
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import pytest
  import numpy
  import torch
  import torch.nn
  import torch.optim
  from torch.utils.data import Dataset, DataLoader
  import bluefog.torch
  from bluefog.common import topology_util
label_desc: Use image node_id 6 for calling a global function with example usage: local_aggregation_train(model, optimizer, dataloader, isCUDA, mini_batch_size) without return types