class_docstr: Applies a caching according to a policy function
code: def backward(
      ctx, inp, models, stored_models, fwd_expert_count, fwd_batch_size
  ):
      (stored_models,) = ctx.saved_tensors
      models, i, fused, num_expert, world_size = ctx.models
      local_experts = models[i]
      models[i] = []  # remove self from list
  
      gradients = [
          [
              _flatten_dense_tensors(
                  [
                      x.grad
                      if torch.is_tensor(x.grad)
                      else torch.zeros(x.shape).cuda()
                      for x in m.parameters()
                  ]
              )
              for m in node
          ]
          for node in models
      ]
      local_gradients = [
          _flatten_dense_tensors(
              [
                  x.grad
                  if torch.is_tensor(x.grad)
                  else torch.zeros(x.shape).cuda()
                  for x in m.parameters()
              ]
          )
          for m in local_experts
      ]
  
      if fused:
          stored_models = stored_models.any(dim=1)
          num_expert = 1
  
      fmoe_cuda.gradient_exchange(
          stored_models,
          local_gradients,
          gradients,
          num_expert,
          world_size,
      )
  
      _update_local_model_params(local_experts, local_gradients)
  
      return (
          inp,
          None,
          None,
          None,
          None,
          None,
          None,
          None,
          None,
          None,
          None,
          None,
          None,
          None,
      )
  
imports: import torch
  from torch.autograd import Function
  import fmoe_cuda
  from .utils import get_torch_default_comm
  from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
  from copy import deepcopy
label_desc: Use image node_id 2 for calling the MOECache obj's underlying member method code with example usage: obj.backward(ctx, inp, models, stored_models, fwd_expert_count, fwd_batch_size) and returns: inp,  None,  None,  None,  None,  None,  None,  None,  None,  None,  None,  None,  None,  None