code: def gradient_tracking(
      X, y, w_opt, loss, maxite=2000, alpha=1e-1, **kwargs
  ):
      if loss == "logistic_regression":
          rho = kwargs.get("rho", 1e-1)
      elif loss == "linear_regression":
          rho = 0
      else:
          raise NotImplementedError(
              "Task not supported. This example only supports"
              + " linear_regression and logistic_regression"
          )
  
      w = torch.zeros(n, 1, dtype=torch.double, requires_grad=True)
      loss_step(
          X,
          y,
          w,
          tensor_name="neighbor.allreduce.Grad.Tracking.w",
          loss=loss,
          rho=rho,
      )
      q = w.grad.data.clone()  # q^0 = grad(w^0)
      w.grad.data.zero_()
  
      grad_prev = q.clone()
      mse = []
      for _ in range(maxite):
          # Algorithm:
          # w^{k+1} = neighbor_allreduce(w^k) - alpha*q^k
          # q^{k+1} = neighbor_allreduce(q^k) + grad(w^{k+1}) - grad(w^k)
  
          # Notice the communication of neighbor_allreduce can overlap with gradient computation.
          w_handle = bf.neighbor_allreduce_nonblocking(
              w.data, name="Grad.Tracking.w"
          )
          q_handle = bf.neighbor_allreduce_nonblocking(
              q, name="Grad.Tracking.q"
          )
          w.data = bf.synchronize(w_handle) - alpha * q
          # calculate local gradient
          loss_step(
              X,
              y,
              w,
              tensor_name="neighbor.allreduce.Grad.Tracking.w",
              loss=loss,
              rho=rho,
          )
          grad = w.grad.data.clone()
          q = bf.synchronize(q_handle) + grad - grad_prev
          grad_prev = grad
          w.grad.data.zero_()
  
          # record convergence
          if bf.rank() == 0:
              mse.append(torch.norm(w.data - w_opt.data, p=2))
  
      return w, mse
  
file_path: bluefog/examples/pytorch_optimization.py
imports: import os
  import torch
  import matplotlib
  import matplotlib.pyplot
  import argparse
  import bluefog.torch
  from bluefog.common import topology_util
  import networkx
label_desc: Use image node_id 8 for calling a global function with example usage: gradient_tracking(X, y, w_opt, loss, maxite, alpha) and returns: w,  mse