[PyTorch] Modify `data_parallel` to work with small tensors (#37704)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/37704
If input tensor can not be chunked, run `parallel_apply` on fewer devices
Modfy input tensor dimention in `DataParallelUsesAllAvailableCUDADevices_CUDA` to be chunkable by any number of available CUDA devices
Test Plan: Run `test/cpp/api/parallel` on machine with 6 GPUs
Differential Revision: D21365416
fbshipit-source-id: 60fdfed4a0e6256b2c966c2ea3e8d0bfb298d9a8
