Solved by driazati in post #2 
              
              
                This is close, the issue is that when self.apply_batchnorm is false, there is no bn attribute on the module, so it cannot be accessed / checked. So 
        if self.apply_batchnorm:
            self.bn = nn.BatchNorm2d(out_channels)

turns into 
        if self.apply_batchnorm:
            self.bn =…