I came across a different solution which uses backward function. It’s all about playing with backward's parameters. More information can be found here. I am still looking for a better solution, if it exists.
1 Like