Solved by ptrblck in post #2 
              
              
                I’ve reimplemented a toy example from this blog post a while ago and you can find the code here. 
Since this code is quite old by now, you might need to change some details (e.g. swap data[0] for .item()).