Hi,
If you want to save it for the backward pass, then you can save it within the context as ctx.i_max = i.max().
If you want to save it for external use, then you can use any pythonic way you want: use a global variable, save it in a list, have a stats object where you can record what happens etc.