code: def _get_ref_results(self, model, input, my_lr):
      with torch.no_grad():
          # Compute one iteration local output.
          weight = model.weight.T.clone().to(self.rank)
          v = torch.Tensor(input[self.rank]).to(self.rank)
          ref_forward_output_my_rank = torch.matmul(v, weight)
          # Compute one iteration global weight update.
          v = torch.Tensor(input[: self.world_size]).to(self.rank)
          grad = (
              v.float()
              .sum(0)
              .repeat(weight.shape[0], 1)
              .div(self.world_size)
          )
          ref_weight_out = weight - grad.T * my_lr
  
      return ref_forward_output_my_rank, ref_weight_out
  
file_path: pytorch/test/distributed/fsdp/test_fsdp_uneven.py
imports: import sys
  import torch
  from torch import distributed
  from torch.distributed.fsdp import FullyShardedDataParallel
  from torch.nn import Linear
  from torch.optim import SGD
  from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
  from torch.testing._internal.common_fsdp import FSDPTest
  from torch.testing._internal.common_utils import run_tests, TEST_WITH_DEV_DBG_ASAN
label_desc: Use image node_id 1 for calling the TestUnevenParamShard obj's underlying member method code with example usage: obj._get_ref_results(model, input, my_lr) and returns: ref_forward_output_my_rank,  ref_weight_out