torch.autograd.function.FunctionCtx.set_materialize_grads
FunctionCtx.set_materialize_grads(value)
Sets whether to materialize grad tensors. Default is "True".
This should be called only from inside the "forward()"
   method
If "True", undefined grad tensors will be expanded to tensors full
   of zeros prior to calling the "backward()" and "jvp()" methods.
Example::
      >>> class SimpleFunc(Function):
      >>>     @staticmethod
      >>>     def forward(ctx, x):
      >>>         return x.clone(), x.clone()
      >>>
      >>>     @staticmethod
      >>>     @once_differentiable
      >>>     def backward(ctx, g1, g2):
      >>>         return g1 + g2  # No check for None necessary
      >>>
      >>> # We modify SimpleFunc to handle non-materialized grad outputs
      >>> class Func(Function):
      >>>     @staticmethod
      >>>     def forward(ctx, x):
      >>>         ctx.set_materialize_grads(False)