arguments to the "run_fn". However, the logic has no way to
  anticipate if the user will move Tensors to a new device within the
  "run_fn" itself.  Therefore, if you move Tensors to a new device
  ("new" meaning not belonging to the set of [current device + devices
  of Tensor arguments]) within "run_fn", deterministic output compared
  to non-checkpointed passes is never guaranteed.
torch.utils.checkpoint.checkpoint(function, args, use_reentrant=True, *kwargs)
Checkpoint a model or part of the model
Checkpointing works by trading compute for memory. Rather than
   storing all intermediate activations of the entire computation
   graph for computing backward, the checkpointed part does not
   save intermediate activations, and instead recomputes them in
   backward pass. It can be applied on any part of a model.
Specifically, in the forward pass, "function" will run in
   "torch.no_grad()" manner, i.e., not storing the intermediate