create_feature_extractor does not track 'if' in forward

🐛 Describe the bug

It looks like create_feature_extractor does not track if statement in forward, for example with the efficientnet architecture:

import torchvision

efficientnetb0 = torchvision.models.efficientnet_b0()
return_nodes = {f'features.{k}': str(v) for v, k in enumerate([3, 4, 6, 7])}

backbone = torchvision.models.feature_extraction.create_feature_extractor(efficientnetb0, return_nodes=return_nodes)

# print first Sequential module for efficientnetb0
print(efficientnetb0.features[1])
>>> Sequential(
  (0): MBConv(
    (block): Sequential(
      (0): ConvNormActivation(
        (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)
        (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (2): SiLU(inplace=True)
      )
      (1): SqueezeExcitation(
        (avgpool): AdaptiveAvgPool2d(output_size=1)
        (fc1): Conv2d(32, 8, kernel_size=(1, 1), stride=(1, 1))
        (fc2): Conv2d(8, 32, kernel_size=(1, 1), stride=(1, 1))
        (activation): SiLU(inplace=True)
        (scale_activation): Sigmoid()
      )
      (2): ConvNormActivation(
        (0): Conv2d(32, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (stochastic_depth): StochasticDepth(p=0.0, mode=row)
  )
)

# print same module for backbone
>>> (1): Module(
    (0): Module(
      (block): Module(
        (0): Module(
          (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)
          (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): SiLU(inplace=True)
        )
        (1): Module(
          (avgpool): AdaptiveAvgPool2d(output_size=1)
          (fc1): Conv2d(32, 8, kernel_size=(1, 1), stride=(1, 1))
          (activation): SiLU(inplace=True)
          (fc2): Conv2d(8, 32, kernel_size=(1, 1), stride=(1, 1))
          (scale_activation): Sigmoid()
        )
        (2): Module(
          (0): Conv2d(32, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)
          (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
      )
    )
  )

Is the StochasticDepth layer implemented but hidden in the backbone or is it missing (due to the if) ?

Versions

Collecting environment information...
PyTorch version: 1.10.0
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: CentOS Linux release 8.2.2004 (Core) (x86_64)
GCC version: (GCC) 8.3.1 20191121 (Red Hat 8.3.1-5)
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.9

Python version: 3.6.9 |Anaconda, Inc.| (default, Jul 30 2019, 19:07:31) [GCC 7.3.0] (64-bit runtime)
Python platform: Linux-4.18.0-193.6.3.el8_2.x86_64-x86_64-with-centos-8.2.2004-Core
Is CUDA available: False
CUDA runtime version: Could not collect
GPU models and configuration: No devices found.
Nvidia driver version: Could not collect
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.1
[pip3] torch==1.10.0
[pip3] torchvision==0.11.1
[conda] blas 1.0 mkl
[conda] cudatoolkit 10.2.89 hfd86e86_1
[conda] ffmpeg 4.3 hf484d3e_0 pytorch
[conda] mkl 2020.2 256
[conda] mkl-service 2.3.0 py36he8ac12f_0
[conda] mkl_fft 1.3.0 py36h54f3939_0
[conda] mkl_random 1.1.1 py36h0573a6f_0
[conda] numpy 1.19.1 py36hbc911f0_0
[conda] numpy-base 1.19.1 py36hfa32c7d_0
[conda] pytorch 1.10.0 py3.6_cuda10.2_cudnn7.6.5_0 pytorch
[conda] pytorch-mutex 1.0 cuda pytorch
[conda] torchvision 0.11.1 py36_cu102 pytorch

Start with the create_feature_extractor entry point and run the EfficientNet reproducer from the issue. Trace the MBConv forward path around StochasticDepth and determine whether the layer is intentionally omitted or lost; done means the behavior is explained and covered by a regression test or a clearly documented expected result.