@cbaziotis, Can you explain a bit about the lengths parameter in forward()? I was looking through your gist https://gist.github.com/cbaziotis/94e53bdd6e4852756e0395560ff38aa4#file-selfattention-py-L31 109. It looks like that parameter is a Variable, but I am not sure where it comes from.