setting up the "ctx" object. "output" is the output of the
     forward, "inputs" are a Tuple of inputs to the forward.

See Extending torch.autograd for more details

The context can be used to store arbitrary data that can be then
   retrieved during the backward pass. Tensors should not be stored
   directly on ctx (though this is not currently enforced for
   backward compatibility). Instead, tensors should be saved either
   with "ctx.save_for_backward()" if they are intended to be used in
   "backward" (equivalently, "vjp") or "ctx.save_for_forward()" if
   they are intended to be used for in "jvp".
Return type:
      Any