SplitParam
Attribute |
Type |
Default Value |
Description |
|---|---|---|---|
split_dim |
int |
0 |
- |
split_num |
int |
2 |
- |
split_sizes |
List[int] |
list() |
- |
Restrictions
The length of split_sizes must be less than that of DEFAULT_SVECTOR_SIZE. Otherwise, a length overflow exception is thrown.
Example
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 | import torch import torch_atb def split(): input_npu = torch.randn(6, 6, dtype=torch.float16).npu() print("input: ", input_npu) split_param = torch_atb.SplitParam(split_dim = 0, split_num = 2) split = torch_atb.Operation(split_param) def split_run(): split_outputs = split.forward([input_npu]) return split_outputs outputs = split_run() print("outputs: ", outputs) if __name__ == "__main__": split() |
Parent topic: OpParam