examples/demo_jax_distributed.py/predict
def predict(data):
    predictions, updated_non_trainable_variables = model.stateless_call(
        device_train_state.trainable_variables,
        device_train_state.non_trainable_variables,
        data,
    )
    return predictions