The leak is likely because the @weak_script decorator captures the frame information in which the function is executing, inadvertently also capturing tensors. Since the weak script functions are never compiled, it then never frees the frame and the tensors leak.

It also seems like those inner functions are not really closures in the first place. Can't we just move them out of the function, which would avoid the leak?

It also seems like those inner functions are not really closures in the first place. Can't we just move them out of the function, which would avoid the leak?

Thanks for the explanation. I feel it's worth more investigations. I could create a benchmark case to reproduce the problem or even report an issue for this.

As obvious, we don't expect and want anyone to actually call/use those inner functions. We feel safe to leave them inside.

@zhangguanheng66 merged this pull request in 3caf4e6.

@zdevito - the functions are very useful abstractions for multihead attention, but we don't want to make them available at a functional.py level.