Changed some code to let it work