Models fprop should support only one structure. The caller of this function should be responsible for converting the input_batch to the right structure.
Models fprop should support only one structure. The caller of this function should be responsible for converting the input_batch to the right structure.

PiperOrigin-RevId: 458253257

google/praxis,google/praxis