Update grid.py

I have made it really simple and put it inside a function