Fix PyType errors in input normalization wrapper.
Fix PyType errors in input normalization wrapper.

PiperOrigin-RevId: 389885508
Change-Id: I4b6d5984e933a61a64ce92732c9a81b3087f2058

deepmind/acme,deepmind/acme