[solved] Tensor groups in forward function

[kaiyuyue]

torch.split and torch.cat will be ok. Thanks.