I think I'll stick with the original idea of just accelerating the tensor operations, because that will be less trouble of bringing it together with PyTorch tho 🤔