code: def _compute(module, *inputs, do_checkpoint=False):
      if do_checkpoint:
          outputs = ckpt(module, *inputs)
      else:
          outputs = module(*inputs)
  
      if torch.is_tensor(outputs):
          outputs = (outputs,)
  
      sum(
          o.sum()
          for o in outputs
          if torch.is_tensor(o) and o.requires_grad
      ).backward()
  
      grads = [p.grad for p in module.parameters()]
      input_grads = [inp.grad for inp in inputs if torch.is_tensor(inp)]
  
      return {
          "outputs": outputs,
          "module_grads": grads,
          "input_grads": input_grads,
      }
  
imports: import pytest
  import torch
  import deepspeed
  from deepspeed.accelerator import get_accelerator
  from copy import deepcopy
  from unit.common import DistributedTest
label_desc: Use image node_id 1 for calling a global function with example usage: _compute(module) and returns: dict