vscode: evaluating python code snippets generates strange errors

Question

I am trying to run code snippets in the latest vscode 1.66 with the latest python extension. I tried to evaluate a function as listed below--for a reinforcement learning experiment--but I am getting strange errors about the return being outside of function--even though that is not the case.

a = 1
def random_rollout(env):
    state = env.reset()
    done = False
    cumulative_reward = 0.0
    while not done:
        action = np.random.choice([0, 1])
        state, reward, done, _ = env.step(action)
        cumulative_reward += reward
    return cumulative_reward

reward = random_rollout(env)

Now I read about vscode having issues with evaluating code blocks because the python extension has errors with indentation. I think that the indentation is just fine in the code block, because if I run this function from the terminal I get no errors. But there were some strange tricks that should let me get around this problem right--like making sure I evaluate a block that start with an unindented block and ends with an unindented block, etc.

I have been reading all over google, but I am more confused now than when I started. So what would be the right way to evaluate this code block in vscode and python, using the SHIFT-Enter key.