I don’t think there are hard rules about it, so I would look at the torch.autograd._functions for inspiration when to use which and see what looks similar to your case.
For example in the Clamp Function 232 you store the mask in the ctx while for the square root 104 a few lines down, you just store the input.
Best regards
Thomas