Missing test_img and train_img

[MakuSimz]

these lists are not defined in your code, i do not think this code works as of yet

[hlamba28]

Thanks for pointing it out. Required changes done, closing the issue.