If n_samples * n_features < 1e8, and you have a medium to good GPU (>=GTX1080), then there is no need to split by class, below is the “vectorized” implementation, not tested.
# new implementation of kronecker
def kronecker(A, B):
    return torch.einsum('nab,ncd->nacbd', A, B)\
        .view(A.size(0), A.size(1)*B.size(1), A.size(2)*B.size(2))

Main code, no need to use pools or whatever:
# according to your code, the shape of `X_prime` should be (n_samples, n_features + 1)
# whether encoding = "amplit" or "stereo"
rows = n_samples
cols = n_features + 1


# you may keep this if "n_samples * n_features" > 1e8
# X_prime_class = X_prime[y_class_index == i]

# Number of rows/columns in density matrix
density_nrow_ncol = (n + 1)**self.n_copies

# Encode vectors into quantum densities
# density: [rows, cols, cols], each density[k] is the original `density_each_row`
density = torch.matmul(X_prime.view(rows, cols, 1),
                       X_prime.view(rows, 1, cols))

# Calculate n-fold Kronecker tensor product
if self.n_copies == 1:
    density = density
else:
    density_copy = density
    for u in range(self.n_copies - 1):
        density = kronecker(density, density_copy)

# Calculate sum of quantum densities belonging to either class, per subset split
density_sum = density.sum(dim = 0)

# calculate centroid and q_hels_obs_terms
centroid = (1 / (m_class + 1e-6))*density_sum_class

if self.class_wgt == 'equi':
    q_hels_obs_terms = 0.5*centroid_class
elif self.class_wgt == 'weighted':
    q_hels_obs_terms = (m_class / m)*centroid_class
else:
    raise ValueError('class_wgt should be "equi" or "weighted"')