[docs] improve torch.stack example code to be reproducible (#133857)

Improve the sample code can produce the expected results after copying and executing it.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/133857
Approved by: https://github.com/soulitzer
This commit is contained in:
ooooo
2024-08-21 14:07:02 +00:00
committed by PyTorch MergeBot
parent 585c049fa3
commit 32e052e468

View File

@ -1696,22 +1696,21 @@ Example::
>>> x
tensor([[ 0.3367, 0.1288, 0.2345],
[ 0.2303, -1.1229, -0.1863]])
>>> x = torch.stack((x, x)) # same as torch.stack((x, x), dim=0)
>>> x
>>> torch.stack((x, x)) # same as torch.stack((x, x), dim=0)
tensor([[[ 0.3367, 0.1288, 0.2345],
[ 0.2303, -1.1229, -0.1863]],
[[ 0.3367, 0.1288, 0.2345],
[ 0.2303, -1.1229, -0.1863]]])
>>> x.size()
>>> torch.stack((x, x)).size()
torch.Size([2, 2, 3])
>>> x = torch.stack((x, x), dim=1)
>>> torch.stack((x, x), dim=1)
tensor([[[ 0.3367, 0.1288, 0.2345],
[ 0.3367, 0.1288, 0.2345]],
[[ 0.2303, -1.1229, -0.1863],
[ 0.2303, -1.1229, -0.1863]]])
>>> x = torch.stack((x, x), dim=2)
>>> torch.stack((x, x), dim=2)
tensor([[[ 0.3367, 0.3367],
[ 0.1288, 0.1288],
[ 0.2345, 0.2345]],
@ -1719,7 +1718,7 @@ Example::
[[ 0.2303, 0.2303],
[-1.1229, -1.1229],
[-0.1863, -0.1863]]])
>>> x = torch.stack((x, x), dim=-1)
>>> torch.stack((x, x), dim=-1)
tensor([[[ 0.3367, 0.3367],
[ 0.1288, 0.1288],
[ 0.2345, 0.2345]],