Moving while in the air

I am using a slightly edited version of the Move script found within the Move reference page:

using UnityEngine;
using System.Collections;
 
public class Movement : MonoBehaviour {
	
    public float speed = 15.0F;
    public float jumpSpeed = 12.0F;
    public float gravity = 30.0F;
	
	
		
    private Vector3 moveDirection = Vector3.zero; // (0,0,0)
    
	void Update() {
       
		CharacterController controller = GetComponent<CharacterController>();
       
		if (controller.isGrounded) {
           
			moveDirection = new Vector3(Input.GetAxis("Horizontal"), Input.GetAxis("Vertical"));
           // moveDirection = transform.TransformDirection(moveDirection);
            moveDirection *= speed;
           
			if (Input.GetButton("Jump")){
                
			moveDirection.y = jumpSpeed; 
				
				
			}
 
        }
        moveDirection.y -= gravity * Time.deltaTime;
        controller.Move(moveDirection * Time.deltaTime); 
		// WIP -----------------------------------
		if(moveDirection.y == -20){ 
			
			Debug.Log("Breach on the y!"); 
			
		}
    }

At the moment this enables me to jump and move along the x and y axis, perfect. Unfortunately when I jump the axis are disabled until controller.isGrounded turns true again. I have tried inserting another if statement to deal with the issue but I’m having a lot of trouble. I was wondering what the simplest addition is that I could make to the script so that the player has control of the character while in the air?

Just delete lines 18 and 31.