diff --git a/torchvision/ops/stochastic_depth.py b/torchvision/ops/stochastic_depth.py
index 0b95e7cca67..9ba6394326b 100644
--- a/torchvision/ops/stochastic_depth.py
+++ b/torchvision/ops/stochastic_depth.py
@@ -28,9 +28,10 @@ def stochastic_depth(input: Tensor, p: float, mode: str, training: bool = True)
         return input
 
     survival_rate = 1.0 - p
-    size = [1] * input.ndim
     if mode == "row":
-        size[0] = input.shape[0]
+        size = [input.shape[0]] + [1] * (input.ndim - 1)
+    else:
+        size = [1] * input.ndim
     noise = torch.empty(size, dtype=input.dtype, device=input.device)
     noise = noise.bernoulli_(survival_rate).div_(survival_rate)
     return input * noise