search

snap-stanford / ogb

Benchmark datasets, data loaders, and evaluators for graph machine learning
https://ogb.stanford.edu
MIT License
1.93k stars 397 forks source link

GPU memory leakage in the baseline code of wikikg90m #139

Closed jianyucai closed 3 years ago

jianyucai commented 3 years ago

Hi, thanks for your code.

I noticed that there is a GPU memory leakage problem in the baseline code of wikikg90m.

https://github.com/snap-stanford/ogb/blob/4db49c17a0c51352990b6468803ef5eb495817dd/examples/lsc/wikikg90m/dgl-ke-ogb-lsc/python/dglke/train_pytorch.py#L201-L202

In the above line, the log variable is stored in GPU, but for each sample in validation or testing, it is appended to the logs variable (a list).

I found that this issue could be fixed by the following code (using function .cpu()).

log = model.forward_test_wikikg(query, ans, candidate, sampler.mode, gpu_id).cpu()
hyren commented 3 years ago

Thank you for catching this! We have pushed the fix.