Some changes in the code..