text: import torch.nn as nn
  from torch.nn import init
  import torch.nn.functional as F
  import torchvision.models as models
  
  import nsml
  
  def weights_init_kaiming(m):
      classname = m.__class__.__name__
      if classname.find('Conv') != -1:
          init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
          if m.bias is not None:
              init.constant_(m.bias.data, 0.0)
      elif classname.find('Linear') != -1:
          init.kaiming_normal_(m.weight.data, a=0, mode='fan_out')
          if m.bias is not None:
              init.constant_(m.bias.data, 0.0)
      elif classname.find('BatchNorm1d') != -1:
          init.normal_(m.weight.data, 1.0, 0.02)
          if m.bias is not None:
              init.constant_(m.bias.data, 0.0)
      elif classname.find('BatchNorm2d') != -1:
          init.normal_(m.weight.data, 1.0, 0.02)
          if m.bias is not None:
              init.constant_(m.bias.data, 0.0)
      elif classname.find('BatchNorm') != -1:
          if m.affine is not None:
              init.constant_(m.weight.data, 1.0)
              init.constant_(m.bias.data, 0.0)
  
  
  
  def weights_init_classifier(m):
      classname = m.__class__.__name__
      if classname.find('Linear') != -1:
          init.normal_(m.weight.data, std=0.001)
          init.constant_(m.bias.data, 0.0)
  
  
  class ClassBlock(nn.Module):
      def __init__(self, input_dim, class_num, dropout=True, relu=True, num_bottleneck=512): #512
          super(ClassBlock, self).__init__()
          add_block = []
          add_block += [nn.Linear(input_dim, num_bottleneck)]
          add_block += [nn.BatchNorm1d(num_bottleneck)]
          if relu:
              add_block += [nn.ReLU()]
          if dropout:
              add_block += [nn.Dropout(p=0.5)]
          add_block = nn.Sequential(*add_block)
          add_block.apply(weights_init_kaiming)
          classifier = []
          classifier += [nn.Linear(num_bottleneck, class_num)]
          classifier = nn.Sequential(*classifier)
          classifier.apply(weights_init_classifier)
          self.add_block = add_block
          self.classifier = classifier
      def forward(self, x):
          x = self.add_block(x)
          x = self.classifier(x)
          return x
  
  
  class ResNetSimCLR(nn.Module):
  
      def __init__(self, base_model, out_dim):
          super(ResNetSimCLR, self).__init__()
          self.resnet_dict = {"resnet18": models.resnet18(pretrained=False),
                              "resnet50": models.resnet50(pretrained=False)}
  
          self.resnet = self._get_basemodel(base_model)
          num_ftrs = self.resnet.fc.in_features
  
          self.features = nn.Sequential(*list(self.resnet.children())[:-1])
  
          # projection MLP
          self.l1 = nn.Linear(num_ftrs, num_ftrs)
          self.l2 = nn.Linear(num_ftrs, out_dim)
  
          self.classifier = ClassBlock(2048, 265)
          self.classifier.apply(weights_init_classifier)
  
      def _get_basemodel(self, model_name):
          try:
              model = self.resnet_dict[model_name]
              print("Feature extractor:", model_name)
              return model
          except:
              raise ("Invalid model name. Check the config file and pass one of: resnet18 or resnet50")
  
      def forward(self, x):
          h = self.features(x)
          h = h.squeeze()
  
          x = self.l1(h)
          x = F.relu(x)
          x = self.l2(x)
  
          h = self.classifier(h)
  
          return x, h
  