Hi,
You’re missing k in save_for_backward 
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.