Thanks for the reply, I’ll explain a bit more : let’s say I have a function such as
class MyRound(torch.autograd.Function):
    @staticmethod
    def forward(ctx,input):
        return torch.round(input)
    
    @staticmethod
    def backward(ctx,grad):
        return grad.clone()

In my use case the grad at backward is not well defined. Rather than operating on this badly defined grad I can take a look at the model loss (e.g. cross entropy) and redefine a grad that is useful. Can I access the cross-entropy loss (or whatever “root” loss from which the backpropagation started) from within backward ? (e.g. through the ctx object) Result would look like :
@staticmethod
    def backward(ctx,grad):
        good_grad = ctx.root_loss 
        # and some operations to have the correct grad
        return good_grad