Fixed bug in context switching when a function ends (infinite loop)