Ejecute Algorithm completamente en GPU

¿Es posible ejecutar el siguiente algoritmo completamente en la GPU, usando RAPIDS o CuPy, para mejorar su tiempo de ejecución? El algoritmo es una implementación de https://ieeexplore.ieee.org/document/1407871. En la actualidad se utiliza la implementación de K-Means de la biblioteca cuml, que ha mejorado mucho su rendimiento.

import numpy as np
from cuml import KMeans

def _calc_scores_and_labels(x,k, beta, eps=1e-3, max_iterations=100):
    k_means = KMeans(n_clusters=k)

    # Assign random weights to each feature. Sum of weights should be equal to 1
    random_weights = np.random.rand(x.shape[1])
    exp_weights = np.exp(random_weights)
    weights = exp_weights / np.sum(exp_weights)

    # Compute clusters using samples scaled by weight in the power of beta
    modified_x = x * weights ** beta
    clusters = k_means.fit_predict(modified_x)

    # Compute score for clustering
    old_score = _calc_objective(k_means, modified_x)
    for i in range(max_iterations):

        # Find position of centroids
        divider = weights.copy() ** beta
        divider[abs(divider) < eps] = 1
        centroids = k_means.cluster_centers_ / divider

        # Compute D-array used for computation of new weights
        d = np.zeros(x.shape[1])
        for k in range(d.size):
            d[k] = 0
            for j in range(x.shape[0]):
                d[k] += abs(x[j][k] - centroids[clusters[j]][k])

        # Compute new weights using D-array
        new_weights = np.zeros(weights.size)
        if beta == 1:
            new_weights[np.argmin(d)] = 1
        else:
            for k in range(new_weights.size):
                if abs(d[k]) < eps:
                    continue
                for current_d in d:
                    if abs(current_d) < eps:
                        continue
                    new_weights[k] += (d[k] / current_d) ** (1 / (beta - 1))
                new_weights[k] = 1 / new_weights[k]
        weights = new_weights

        # Recompute clusters and check convergence of algorithm
        modified_x = x * weights ** beta
        clusters = k_means.fit_predict(modified_x)
        new_score = _calc_objective(k_means, modified_x)
        if abs(new_score - old_score) < eps:
            break
        old_score = new_score
    return weights

def _calc_objective(k_means, x):
     return -k_means.score(x)

data=np.random.random((1000,4))
w=_calc_scores_and_labels(data,k=2, beta=1.5)
print(w)