Small change

- Cleaned up code in work section