Resolving tracing problem on StochasticDepth iterator. (#4372)

Co-authored-by: Francisco Massa <fvsmassa@gmail.com>

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