Not really sure what you mean by ‘small code sample’. So I will try:
class Net(nn.Module):
    in_planes = 2048

    def __init__(self, num_classes, model_path,  model_name):
        super(Net, self).__init__()
        self.base = ResNet(block=Bottleneck,
                           layers=[3, 4, 6, 3])

        self.gap = nn.AdaptiveAvgPool2d(1)
        self.num_classes = num_classes
        self.bottleneck = nn.BatchNorm1d(self.in_planes)
        self.bottleneck.bias.requires_grad_(False)  # no shift
        self.classifier = nn.Linear(self.in_planes, self.num_classes, bias=False)

    def forward(self, x):

        global_feat = self.gap(self.base(x))  # (b, 2048, 1, 1)
        global_feat = global_feat.view(global_feat.shape[0], -1)  # flatten to (bs, 2048)

        feat = self.bottleneck(global_feat.unsqueeze(1))  ### To allow SyncBatchnorm

        cls_score = self.classifier(feat.squeeze())   ### To adjust for Linear layer input

        return cls_score, global_feat 


def train():
        model.train()
        optimizer.zero_grad()
        img, target = batch
        img = img.to(device)
        target = target.to(device)t
        score, feat = model(img)
        LOSS = loss_fn(score, feat, target)
        LOSS.backward()
        optimizer.step()

I shortened the code as much as possible to get the most important parts I think. The img size (input in train function) is [batch_size, 3, W, H] so standard for images.
EDIT: format and making a little clearer code