Unable to setup test_dataset

[tanay2001]

The TrainDataset function used to generate test_data is broken , it accepts a rel folder as input while no such folder exists in test.
Even TrainInferDataset is never called but has been imported from loader.py , I feel this is a bug please look into it

[brabbit61]

@tanay2001 faced the same error, just use ~1000 examples from training set to test set. It doesn't affect the final accuracy as the loss stops decreasing after ~5 epochs anyway