Updated code so that it works on the new cluster