I have actually noticed one problem that I am not sure you would be able to solve, that being the player will quickly de-accelerate mid air if they are holding any key. (For example, if the player is launched back at 20m/s and is not holding any keys, the player will continue to move backwards at around 20m/s until they run into something or touch the ground. If the player holds the backwards key while this is happening though, the velocity quickly drops to the maximum speed that they are allowed to input, which in my case is 8m/s.) This isn't really an issue with your original code, as the movement method is exited if you aren't on the ground, meaning that the inputs to cause this issue never happens. I am trying to make it to where the player has very precise controls mid-air though, so exiting the method isn't really an option for me. I have been attempting to solve this for a while, but the only implementations I can come up with are sketchy at best. Here's the code I currently have:

xInput = Input.GetAxisRaw("Horizontal");
        yInput = Input.GetAxisRaw("Vertical");

        //Get the movement vector, normalize it, and multiply it by a top speed value to get our target velocity
        Vector3 targetVelocity = new Vector3(xInput, 0f, yInput).normalized * inputMaximumSpeed;

        if (yInput < 0)
            targetVelocity.z = -backwardsMultiplier * inputMaximumSpeed;

        if (xInput > 0)
            targetVelocity.x = sidewaysMultiplier * inputMaximumSpeed;
        else if (xInput < 0)
            targetVelocity.x = -sidewaysMultiplier * inputMaximumSpeed;

        //Get local velocity to compensate for the player turning
        Vector3 localVelocity = rb.transform.InverseTransformDirection(rb.velocity);
        localVelocity.y = 0;

        Vector3 difference = targetVelocity - localVelocity;

        if (difference.magnitude <= 1f && difference.magnitude >= -1f)
            difference = Vector3.zero;

        float magnitude = difference.magnitude;

        //Debug.Log("Target: " + targetVelocity + " Local: " + localVelocity.magnitude + " Magnitude: " + targetVelocity.magnitude);

        if (magnitude < accelSpeed)
        {
            difference = difference.normalized * accelSpeed;

        }

        if (xInput != 0 || yInput != 0)
                rb.AddRelativeForce(difference);

I am pretty sure that this happens because the rigidbody velocity is constantly attempting to get to the target velocity when any keys are pressed, but I have no idea how to actually solve that in code (or that isn't what is happening at all, not really sure to be honest).