Yes save_for_backward is just for input and outputs, it allows to perform additional checks (make sure that you don’t create non-collectable cycles). For intermediary results, you can save them as attribute of the context yes.