Solved by ZimoNitrome in post #4 
              
              
                It still broke gradients in some specific cases. I finally found this thread giving a weird hack that solved the problem: 


FInal code if someone has a similar problem in the future: 
def quantize(image, palette):
    """
    Similar to PIL.Image.quantize() in PyTorch. Built to maintain gradient.
 …