Yes. The output of the localization network is theta. But it seems that there is no code for localization network in your code.