This is my code, I use same code to test it, but only change the GPU id by os.environ['CUDA_VISIBLE_DEVICES'] = '0', RTX 3090 give me error as I mentioned, but TianXP works fun.
def tp(x, nmb_clusters):
device = x.device
x = common_functions.to_numpy(x).astype(np.float32)
n_data, d = x.shape
# faiss implementation of k-means
clus = faiss.Clustering(d, nmb_clusters)
clus.niter = 20
clus.max_points_per_centroid = 10000000
index = faiss.IndexFlatL2(d)
if faiss.get_num_gpus() > 0:
index = faiss.index_cpu_to_all_gpus(index)
# perform the training
clus.train(x, index)
_, idxs = index.search(x, 1)
return torch.tensor([int(n[0]) for n in idxs], dtype=int, device=device)
if __name__ == '__main__':
a = torch.rand(1000, 256, dtype=torch.float32)
tp(a, 5)
Summary
In the same machine, I have two gpus, one is TianXP, another is RTX 3090, the TianXP works fine, but RTX3090, give the error of
Running on:
Interface:
Reproduction instructions
This is my code, I use same code to test it, but only change the GPU id by
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
, RTX 3090 give me error as I mentioned, but TianXP works fun.