Make sure tf.Module is called before assignments are done in CEB.
Make sure tf.Module is called before assignments are done in CEB.

PiperOrigin-RevId: 356895554
Change-Id: Ieb41282a7f1d25177b97f1012906f4771aefdd99

google-research/pisac