Thanks for checking this issue.
I wonder whether the if-else statement is typo or not?

size if name in self.input_embedding_flags else size