@ptrblck Thanks for the response. Some questions just to clarify, I will have to apply the clustering to my whole training dataset (I guess I do not need to do that for the testing dataset) or only to each corresponding batch size points each time?
For creating the clusters and getting the weights for each cluster the procedure is kind of straight forward:

import torch
import numpy as np
from kmeans_pytorch import kmeans
data_size, dims, num_clusters = 7000, 1, 10 # this is for one sample (set of values/points to be regressed) with 7000 points, my batch size can vary from 4-12 samples while my whole training dataset is 400 samples
x = np.random.randn(data_size, dims)
x = torch.from_numpy(x)
cluster_ids_x, cluster_centers = kmeans(X=x, num_clusters=num_clusters, distance=‘euclidean’, device=torch.device(‘cuda:0’))
classes, class_counts = np.unique(cluster_ids_x, return_counts=True)
weight = 1. / class_counts

Now do I need to assign these weights to each value separately or what? I do not understand that part, and also when you say:
“After creating the weights, you could write a function, which accepts the current output batch with the regression prediction, as well as your cluster centers (k-means dict), and returns a batch of weights, which can then be multiplied to create the final loss.”
This is gonna be in my “CustomLoss()” function, right?