Here's all my Python/Numba (CUDA) code for the encoder block I made :) - View it on GitHub