[BE] fix compute_global_tensor_shape test (pytorch#161441)

Fixes pytorch#161154

**Test**
`pytest  test/distributed/tensor/test_utils.py -s -k test_compute_global_tensor_shape_1D`

Pull Request resolved: pytorch#161441
Approved by: https://github.com/kwen2501

Author: XilunWu

.ci/pytorch/multigpu-test.sh

@@ -45,6 +45,7 @@ if [[ "${SHARD_NUMBER:-2}" == "2" ]]; then
     # DTensor tests
     time python test/run_test.py --verbose -i distributed/tensor/test_random_ops
     time python test/run_test.py --verbose -i distributed/tensor/test_dtensor_compile
+    time python test/run_test.py --verbose -i distributed/tensor/test_utils.py
 
     # DeviceMesh test
     time python test/run_test.py --verbose -i distributed/test_device_mesh

torch/distributed/tensor/_utils.py

@@ -284,12 +284,12 @@ def compute_global_tensor_shape(
     if isinstance(placements[0], Replicate):
         return shape
     elif isinstance(placements[0], Shard):
-        local_shape = torch.tensor(list(shape))
+        local_shape = torch.tensor(list(shape), device=mesh.device_type)
         gathered_shaped_tensors = [
             torch.empty_like(local_shape, device=local_shape.device)
             for _ in range(mesh.size())
         ]
-        funcol.all_gather_inplace(gathered_shaped_tensors, local_shape)
+        funcol.all_gather_inplace(gathered_shaped_tensors, local_shape, mesh)
         sharded_dim_sum = 0
         shard_dim = placements[0].dim
         other_dims = [d for d in range(mesh.ndim) if d != shard_dim]