[BE] Use allocator to allocate workspace (#107178)
As suggested in https://github.com/pytorch/pytorch/pull/106844#discussion_r1293839247 it's better to just allocate DataPtr than the whole tensor
Pull Request resolved: https://github.com/pytorch/pytorch/pull/107178
Approved by: https://github.com/albanD
ghstack dependencies: #106977, #106844
