Revert "[dtensor] move pad/unpad_tensor to separate utils (#124871)"

This reverts commit 0b0eea222978e6b377e2c67f89902d5eb1aa7da3.

Reverted https://github.com/pytorch/pytorch/pull/124871 on behalf of https://github.com/jeanschmidt due to Broke internal tests, see D56587991 for more details ([comment](https://github.com/pytorch/pytorch/pull/124871#issuecomment-2079001103))
This commit is contained in:
PyTorch MergeBot
2024-04-26 09:30:34 +00:00
parent 35a82d4a4a
commit 359ff49bf4
3 changed files with 31 additions and 32 deletions

View File

@ -809,10 +809,8 @@ class TestDTensorPlacementTypes(DTensorTestBase):
]
assert_array_equal(expected_pad_sizes, pad_sizes)
from torch.distributed._tensor._collective_utils import unpad_tensor
unpadded_list = [
unpad_tensor(tensor, shard_placement.dim, pad_sizes[i])
shard_placement._unpad_tensor(tensor, pad_sizes[i])
if pad_sizes[i] > 0
else tensor
for i, tensor in enumerate(splitted_tensor_list)