Custom backward pass

[ptrblck]

Cross-posting from Twitter:
Check out this post and see, if it’s applicable for this use case.