How to do slicing in pytorch? I tried narrow and numpy slicing? Both are not working for outputing train data and test data. Is there any suggestion for solving it?