Adding flexbox for grids