How to convert Tensorflow into PyTorch?