fix torch error "Trying to backward through the graph a second time ..."

[aanjaa]

No description provided.

[saleml]

Thank you for the PR @aanjaa.
Please note this branch of the code will be obsolete in a few days, as we're pushing to publish v1.
As this is your second PR, I'm assuming you are using this code base for your project. Please bear in mind that the API will have significant changes and simplifications (allowing for a wider range of applications). The development branch is called generic-gfn.