Ok, after extracting and applying the weights on the values, I used a fair enough number of 50 clusters over my training dataset values, what I am noticing is that now my regressed values are generally overestimated:

The custom loss that I’ve used is based on MAE:

class L1Weighted(torch.nn.Module):
    def __init__(self):
        super(L1Weighted, self).__init__()

    def forward(self, predicted, true, weights):
        errors = torch.abs(predicted - true)
        return torch.mean(errors*weights)

I am trying to think if there is any other approach that I could try.