Wrap a function to save memory for backward. The returned function performs same
computation as ``func``, but save memory by discarding intermediate results.
It calculates the gradient by recomputing ``func`` using the same input during backward.