Hi! How should I declare an optional BatchNorm submodule in a valid TorchScript way?
class Downsample(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        apply_batchnorm: bool = False,
    ):
        super(Downsample, self).__init__()
        self.conv = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=2,
            bias=False,
            padding=1,
        )
        self.apply_batchnorm = apply_batchnorm
        if self.apply_batchnorm:
            self.bn = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        x = self.conv(x)
        if self.apply_batchnorm:
            x = self.bn(x)
        x = F.leaky_relu(x, 0.3)
        return x

When trying to run torch.jit.script(Downsample(1, 2, 3)) I get following error.
RuntimeError: 
Module 'Downsample' has no attribute 'bn' :
  File "path/to/file.py", line 100
        x = self.conv(x)
        if self.apply_batchnorm:
            x = self.bn(x)
                ~~~~~~~ <--- HERE
        x = F.leaky_relu(x, 0.3)
        return x

I understand that in TorchScript every variable must have single static type. Adding bn: Optional[nn.BatchNorm2d] to the class definition does not help.