diff --git a/python/tutorials/01-vector-add.py b/python/tutorials/01-vector-add.py
index 5ac5a1225..4baa951c1 100644
--- a/python/tutorials/01-vector-add.py
+++ b/python/tutorials/01-vector-add.py
@@ -55,7 +55,7 @@ def add(x: torch.Tensor, y: torch.Tensor):
     # We need to preallocate the output
     output = torch.empty_like(x)
     assert x.is_cuda and y.is_cuda and output.is_cuda
-    n_elements = output.shape[0]
+    n_elements = output.numel()
     # The SPMD launch grid denotes the number of kernel instances that run in parallel.
     # It is analogous to CUDA launch grids. It can be either Tuple[int], or Callable(metaparameters) -> Tuple[int]
     # In this case, we use a 1D grid where the size is the number of blocks