Thank you for your reply InnovArul.
Edit: In case someone else is looking to do something similar the solution seems to be:
class ExpMeanLog(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, running_ema):
        ctx.save_for_backward(input, running_ema)
        return input.exp().mean().log()

    @staticmethod
    def backward(ctx, grad_output):
        input, running_ema = ctx.saved_tensors
        return grad_output * input.exp() / running_ema / input.shape[0], None