Is there a way to split a batch of say N by K where N is the number of samples with K being the feature dimensionality into equal chunks? Torch.split() does the job but it, unfortunately, returns a list which I don’t think you can backprop through. torch.view() could also work but i’m not sure on the order of traversal as I need make sure that each chunk in the new split tensor corresponds to the same order in the unsplit tensor, i.e. if i had a tensor from [1,2,3,4,5,6] and split by two i would like to preserve [[1,2], [3,4], [5,6]].