Added support for python3 and latest versions of PyTorch

[jiteshm17]

Hi @koomri ,

Let me start of with a huge shoutout for the paper and the code. I found it to be immensely helpful. Thanks a lot for your efforts

However, it looks like the code has been written in Python2 and a very old version of PyTorch, both of them which are now almost deprecated. I've created a PR here that contains changes to the code to account for the above details and also bringing about significant speedup in terms of performance. Do take a look at it if you get a chance

For anyone else reading this, I have a working version of this PR forked in my repository which you can find here

Hope this helps someone. Cheers ✌️