Not really an issue, but I made a PyTorch implementation for the discrete model, available here, if anyone is interested.