Thank you very much for your work, I am not very familiar with the tensorflow version, do you have this code for pytorch?