For now I created the following methods:
def distance(color_a, color_b):
    differences = torch.stack([(c_a - c_b) ** 2 for c_a, c_b in zip(color_a, color_b)])
    return torch.sum(differences, dim=-2)**0.5 # Euclidian distance
    
def quantize(image, palette):
    # Assume RGB
    flat_img = image.view(1, 3,-1) # [C, H, W] -> [1, C, H*W]
    img_per_palette_color = torch.cat(len(palette)*[flat_img]) # [1, C, H*W] -> [n_colors, C, H*W]
    distance_per_pixel = distance(img_per_palette_color, palette.unsqueeze(-1)) # [n_colors, C, H*W] -> [n_colors, H*W]
    color_indices = torch.argmin(distance_per_pixel, dim=0) # [n_colors, H*W] -> [H*W]
    new_colors = palette[color_indices].T # [H*W] -> [C, H*W]
    return new_colors.view(image.shape) # [C, H*W] -> [C, H, W]

colors = torch.tensor([
    [0, 75, 135],
    [255, 205, 0]
])
colors_normalized = colors/255

original = TF.to_tensor(Image.open(r"[PATH]").convert("RGB"))
display(TF.to_pil_image(original))

quantized = quantize(original, colors_normalized)
display(TF.to_pil_image(quantized))

Which outputs:
Code_HhfocMSnx7294×593 11.3 KB
(1st one is the original, 2nd one is quantized)
I will investigate if this breaks the gradient and if I can make the function general for multiple images and multiple palettes.