Add support for torch.tensor_split to accept a tensor for `indices` argument (#49169)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/49169
Trying to solve PR request https://github.com/pytorch/pytorch/issues/47479.
This diff tries to overload method `torch.tensor_split` to also accept a tensor for argument `split_size_or_sections` which currently accepts a python list or int. The motivation is to avoid converting a tensor to a list so that when tracing a model/module the tensor operations can be recorded.
Implementation is following the diff that originally added the `tensor_split` method D24166164 (https://github.com/pytorch/pytorch/commit/ef4817fe5a16ba9969562911c5363736a1003bb0).
Test Plan:
```
buck test caffe2/test:torch -- tensor_split
```
https://www.internalfb.com/intern/testinfra/testconsole/testrun/5910974550563805/
```
buck test caffe2/test:others -- tensor_split
```
https://www.internalfb.com/intern/testinfra/testconsole/testrun/1688849905082678/
Reviewed By: mruberry
Differential Revision: D25440885
fbshipit-source-id: 6705dc551279e3a5eb1e5ec1ede2728eab85ffb1
