albanD:

You’re missing k in save_for_backward 


I’m not getting it, sorry…



 albanD:

Also keep in mind that you should use save_for_backward() only for input or output Tensors. Other intermediary Tensors or input/output of other type can just be saved in the ctx as ctx.mat_shape = mat.shape in your case.


Ah, good news, thanks for that!