OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

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

Browse Source 
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
1 changed files with 5 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
torch/_torch_docs.py
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]],