def __init__(self, in_channels, num_classes, elu):
        """
        conv the output down to channels as the desired classesv
        """
        super(OutputTransition, self).__init__()
        self.conv1 = nn.Conv3D(in_channels,
                               num_classes,
                               kernel_size=5,
                               padding=2)
        self.bn1 = nn.BatchNorm3D(num_classes)
        self.relu1 = nn.ELU() if elu else nn.PReLU(num_classes)

        self.conv2 = nn.Conv3D(num_classes, num_classes, kernel_size=1)

        conv the output down to channels as the desired classesv
        
EISeg/med3d/EISegMed3D/inference/models.py