Python第k个素数问题的求解

我写的代码，跑不出来，万般无奈来了这里，希望能够得到解答。万分感谢

import math


# 判断素数
def is_prime(num):
    flag = True
    if num > 1:
        for i in range(2, math.floor(math.sqrt(num)) + 1):
            if (num % i) == 0:
                flag = False
                break
    return flag


# 键盘输入k
k = int(input("请输入k值："))
num = 2
total = 0

while True:
    if is_prime(num):
        total += 1
    if total ==k:
        print(num)
        break
    num += 1