code: import torch
  from torch import nn
  from torch.utils.data import DataLoader, Dataset
  
  
  class DatasetSplit(Dataset):
      def __init__(self, dataset, idxs):
          self.dataset = dataset
          self.idxs = list(idxs)
  
      def __len__(self):
          return len(self.idxs)
  
      def __getitem__(self, item):
          image, label = self.dataset[self.idxs[item]]
          return image, label
  
  
  class LocalUpdate(object):
      def __init__(self, args, dataset=None, idxs=None):
          self.args = args
          self.loss_func = nn.CrossEntropyLoss()
          self.selected_clients = []
          self.ldr_train = DataLoader(DatasetSplit(dataset, idxs), batch_size=self.args.local_bs, shuffle=True, num_workers=4)
  
  
      def train(self, net, history_dict, lrT):
          net.train()
          # train and update
          optimizer = torch.optim.SGD(net.parameters(), lr=lrT, momentum=0.9, weight_decay=5e-4)
  
          epoch_loss = []
          for iter in range(self.args.local_ep):
              batch_loss = []
              for batch_idx, (images, labels) in enumerate(self.ldr_train):
                  images, labels = images.to(self.args.device), labels.to(self.args.device)
                  net.zero_grad()
                  log_probs = net(images)
                  loss = self.loss_func(log_probs, labels)
                  loss.backward()
                  optimizer.step()
                  batch_loss.append(loss.item())
              epoch_loss.append(sum(batch_loss)/len(batch_loss))
          
          current_dict = net.state_dict()
  
          for k in current_dict.keys():
              current_dict[k] -= history_dict[k]
  
          return current_dict, sum(epoch_loss) / len(epoch_loss)
  
  
apis: torch.nn.CrossEntropyLoss