only supports "torch.autograd.backward()" and only if its
     inputs argument is not passed. "torch.autograd.grad()" is not
     supported. If "use_reentrant=False" is specified, checkpointing
     will work with "torch.autograd.grad()".
Parameters:
      * function -- describes what to run in the forward pass of
        the model or part of the model. It should also know how to
        handle the inputs passed as the tuple. For example, in LSTM,
        if user passes "(activation, hidden)", "function" should
        correctly use the first input as "activation" and the second
        input as "hidden"
  * **preserve_rng_state** (*bool**, **optional*) -- Omit stashing
    and restoring the RNG state during each checkpoint. Default:
    "True"

  * **use_reentrant** (*bool**, **optional*) -- Use checkpointing
    implementation that requires re-entrant autograd. If
    "use_reentrant=False" is specified, "checkpoint" will use an
