k = 10
n_init = 10
max_iter = 300
kmeans = faiss.Kmeans(d=data.shape[1], k=k, niter=max_iter, nredo=n_init, gpu=True)
kmeans.train(data.astype(np.float32))

e = time.time()
print("Training time = {}".format(e - s))


s = time.time()

kmeans.index.search(data.astype(np.float32), 1)[1]

e = time.time()
print("Prediction time = {}".format((e - s) / data_size))