PyTorch version of the code

Hello, we are preparing for the work of cvpr2020, I hope you can quote your work, but the edflow library you use, although it runs very well, its debugging and use is still very difficult for me, I still hope to refer to pure pytorch Version of the source code

