Remove unused `train_dataset`

[leosh64]

The dataset train_dataset is unused in the test script. If this line is not removed, the dataset used for testing is required to have a trainA and trainB directory, which might not exist for pure testing datasets. Thus it should be removed.