I have been trying to implement a “color quantization module” for a project but I am getting stuck.
I have tried to simplify the techniques I am using to play around with what is actually possible in PyTorch.
Consider:
n = 10

for input, target in dataloader:
    output = model(input)
    output = (n*output).round() / n   # <---
    loss = l1_loss(output, target)
    loss.backward()

Here, we limit the values that output can take to only 10 unique ones. This makes the values discrete and does not throw any error that the gradient is broken BUT training seems pretty much impossible
Is there any way to train models this way? What I want to do in the future is to vary n per sample.
Maybe torch.quantization has a use case for this? I am not very familiar with that module and have been under the impression that it is done for performance sake and might use a static resolution (n).