Fixed a little bug from Ep3 code