RuntimeError: CUDA out of memory.

[Nireil]

when the points number is more than 50000, the OOM error appears. My GPU memory i s 10.92GiB, is there any way to reduce the memory occupy?

[Nireil]

It seems that memory leak exists in this code

[ShaoxiongYao]

I have the same issue, but it may not be caused by a memory leak.

My code is:

dim = 3

num_p = 80000
num_q = 80000

ref = torch.randn((1, num_p, dim)).float().cuda()
query = torch.randn((1, num_q, dim)).float().cuda()

start_time = time.time()
knn = KNN(1, transpose_mode=True)
d, i = knn(ref, query)
print("query time:", time.time()-start_time)

print("d shape:", d.shape)
print("i shape:", i.shape)

I think this repo computes the distance matrix directly(?), so 80000x80000 cannot be allocated on the GPU. Roughly, assuming each element in the matrix takes one byte, this matrix requires (80000x80000)/(1024*1024*1024) ~ 5.9GB.

[JayKarhade]

@ShaoxiongYao did you find a fix or replacement for this issue?