Best alternative to torch.tensor_split without upgrading to 1.8

[Gabriel_Simmons]

Of course, thanks! Also saw this solution (essentially the same) just using indexing.