Fixed the previous problem ina  better way