Make StochasticDepth FX-compatible (#4373)

test/test_backbone_utils.py

@@ -39,7 +39,7 @@ class TestFxFeatureExtraction:
         'num_classes': 1,
         'pretrained': False
     }
-    leaf_modules = [torchvision.ops.StochasticDepth]
+    leaf_modules = []
 
     def _create_feature_extractor(self, *args, **kwargs):
         """

torchvision/ops/stochastic_depth.py

 import torch
+import torch.fx
 from torch import nn, Tensor
 
 
@@ -37,6 +38,9 @@ def stochastic_depth(input: Tensor, p: float, mode: str, training: bool = True)
     return input * noise
 
 
+torch.fx.wrap('stochastic_depth')
+
+
 class StochasticDepth(nn.Module):
     """
     See :func:`stochastic_depth`.