Skip to content

Conversation

daniellovera
Copy link

This fix enables garbage collection to appropriately work when

def test_net(net, image, text_threshold, link_threshold, low_text, cuda, poly, refine_net=None):
returns, by deleting x we moved to the GPU after we move the forward pass results back to the CPU. Likewise, the same in deleting y and feature after they are no longer needed.

See https://pytorch.org/blog/understanding-gpu-memory-2/#why-doesnt-automatic-garbage-collection-work for more detail.

Running torch.cuda.empty_cache() in test_net() before returning allows nvidia-smi to be accurate.

Relevant package versions
torch version 2.2.1+cu121
torchvision 0.17.1+cu121

Hope this helps!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant