fixed important parts of function in wrapper.py