try this, you simply give the WaitForKeyDown coroutine a list of keys to check, once it registers a key, it will exit the checking loop & terminate:

    using System;
    using System.Collections;
    using System.Collections.Generic;
    using UnityEngine;

    public class InputTester : MonoBehaviour {
        int choice = -1; // valid choices for this example are 0 and 1

        void Start() {
            StartCoroutine(DoSomeLogic());
        }

        IEnumerator DoSomeLogic() {
            print("Starting...");
            print("Awaiting your choice - 0 or 1. :: ");
            yield return StartCoroutine(WaitForKeyDown(new KeyCode[] { KeyCode.Alpha0, KeyCode.Alpha1 }));
            print("Continuing...");
            switch (choice) {
                // logic based on the choice
            }
        }

        IEnumerator WaitForKeyDown(KeyCode[] codes) {
            bool pressed = false;
            while (!pressed) {
                foreach (KeyCode k in codes) {
                    if (Input.GetKey(k)) {
                        pressed = true;
                        SetChoiceTo(k);
                        break;
                    }
                }
                yield return null; //you might want to only do this check once per frame -> yield return new WaitForEndOfFrame();
            }
        }

        private void SetChoiceTo(KeyCode keyCode) {
            switch (keyCode) {
                case (KeyCode.Alpha0):
                    choice = 0;
                    break;
                case (KeyCode.Alpha1):
                    choice = 1;
                    break;
            }
            print(choice);
        }
    }