Thanks for sharing the great code! I am working most of the time with PyTorch. My question is How can I plug my `PyTorch model` to your snake's brain?