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