I have managed to debug this issue and to my surprise this has something to do with gradio and not the original code.