text: #    Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany
  #
  #    Licensed under the Apache License, Version 2.0 (the "License");
  #    you may not use this file except in compliance with the License.
  #    You may obtain a copy of the License at
  #
  #        http://www.apache.org/licenses/LICENSE-2.0
  #
  #    Unless required by applicable law or agreed to in writing, software
  #    distributed under the License is distributed on an "AS IS" BASIS,
  #    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  #    See the License for the specific language governing permissions and
  #    limitations under the License.
  import torch
  from nnunet.network_architecture.generic_UNet import Generic_UNet
  from nnunet.network_architecture.initialization import InitWeights_He
  
  from nnunet.training.network_training.nnUNetTrainerV2 import nnUNetTrainerV2
  from nnunet.utilities.nd_softmax import sigmoid_helper
  from torch import nn
  
  try:
      from torch.nn.functional import gelu
  except ImportError:
      gelu = None
  
  
  class GeLU(nn.Module):
      def __init__(self):
          super().__init__()
          if gelu is None:
              raise ImportError('You need to have at least torch==1.7.0 to use GeLUs')
  
      def forward(self, x):
          return gelu(x)
  
  
  class nnUNetTrainerV2_GeLU(nnUNetTrainerV2):
      def initialize_network(self):
          """
          - momentum 0.99
          - SGD instead of Adam
          - self.lr_scheduler = None because we do poly_lr
          - deep supervision = True
          - ReLU
          - i am sure I forgot something here
  
          Known issue: forgot to set neg_slope=0 in InitWeights_He; should not make a difference though
          :return:
          """
          if self.threeD:
              conv_op = nn.Conv3d
              dropout_op = nn.Dropout3d
              norm_op = nn.InstanceNorm3d
  
          else:
              conv_op = nn.Conv2d
              dropout_op = nn.Dropout2d
              norm_op = nn.InstanceNorm2d
  
          norm_op_kwargs = {'eps': 1e-5, 'affine': True}
          dropout_op_kwargs = {'p': 0, 'inplace': True}
          net_nonlin = GeLU
          net_nonlin_kwargs = {}
          self.network = Generic_UNet(self.num_input_channels, self.base_num_features, self.num_classes,
                                      len(self.net_num_pool_op_kernel_sizes),
                                      self.conv_per_stage, 2, conv_op, norm_op, norm_op_kwargs, dropout_op, dropout_op_kwargs,
                                      net_nonlin, net_nonlin_kwargs, True, False, lambda x: x, InitWeights_He(),
                                      self.net_num_pool_op_kernel_sizes, self.net_conv_kernel_sizes, False, True, True)
          if torch.cuda.is_available():
              self.network.cuda()
          self.network.inference_apply_nonlin = sigmoid_helper
  