Your suggestion works for fc-layer.  But I am wondering: is there a way to compute the jabocian matrix regarding to convolutional layers ?
Flatten both input and output of a conv-layer is not feasible (i.e. [batch, 64, 28, 28] ==> [batch, 64 * 784]), since it would make tensor variable not a leaf node, which has no grad.
Thanks