Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

wrap handle removal in finally #2

Open
amith-ananthram opened this issue Dec 3, 2022 · 0 comments
Open

wrap handle removal in finally #2

amith-ananthram opened this issue Dec 3, 2022 · 0 comments

Comments

@amith-ananthram
Copy link

Hey, thanks for putting this together and sharing it. I have a PR to make this change to torch_intermediate_layer_getter.py:

        try:   
            if self.keep_output:
                output = self._model(*args, **kwargs)
            else:
                self._model(*args, **kwargs)
                output = None
        finally:
            for h in handles:
                h.remove()

In the happy path the code works great but if you try to use it in code that backs off your batch size post OOM, the hooks that are registered with the model's layers retain a reference to the ret tensor dictionary preventing it from being GCd (resulting in less available GPU space). Happy to push my branch with this change for you to review (though I think you might need to enable it somehow?). Thanks!

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

No branches or pull requests

1 participant