Skip to content

wrap handle removal in finally #2

Description

@amith-ananthram

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!

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions