Rewrite kernel.py in Pytorch and add tests
Rewrite kernel.py in Pytorch and add tests