Kunlun Li

@ptrendx I've finished the revision, could you help to find someone to review it?

Added a unit test and signed off my commit.

@timmoon10 Can you take a look and merge this?

> Are we reliant on the user to clear the additional memory? Alternate could to free this memory during the forward pass or something, if the only use here is...

OK, I changed it to use warnings instead of raising an error.