oneflow.vsplit¶
-
oneflow.
vsplit
()¶ Splits input, a tensor with two or more dimensions, into multiple tensors vertically according to indices_or_sections. Each split is a view of input. This is equivalent to calling oneflow.tensor_split(input, indices_or_sections, dim=0) (the split dimension is 0), except that if indices_or_sections is an integer it must evenly divide the split dimension or a runtime error will be thrown. The documentation is referenced from: https://pytorch.org/docs/1.10/generated/torch.vsplit.html.
- Parameters
input (Tensor) – the input tensor.
indices_or_sections (int or a list) – If indices_or_sections is an integer n , input is split into n sections along dimension dim.If input is divisible by n along dimension dim, each section will be of equal size, input.size (dim) / n. If input is not divisible by n, the sizes of the first int(input.size(dim) % n). sections will have size int(input.size(dim) / n) + 1, and the rest will have size int(input.size(dim) / n). If indices_or_sections is a list or tuple of ints, then input is split along dimension dim at each of the indices in the list, tuple or tensor. For instance, indices_or_sections=[2, 3] and dim=0 would result in the tensors input[:2], input[2:3], and input[3:].If indices_or_sections is a tensor, it must be a zero-dimensional or one-dimensional long tensor on the CPU.
- Returns
the output TensorTuple.
- Return type
oneflow.TensorTuple
For example:
>>> import oneflow as flow >>> input = flow.rand(4, 4, 5, 6) >>> output = flow.vsplit(input, (1, 3)) >>> output[0].size() oneflow.Size([1, 4, 5, 6]) >>> output[1].size() oneflow.Size([2, 4, 5, 6]) >>> output[2].size() oneflow.Size([1, 4, 5, 6])