Update the train script to reduce the memory further by clearing train features variable at the expense of loading time