Model I am using is
class Net(nn.Module):
    """backbone + projection head"""
    def __init__(self, name='resnet50', head='mlp', feat_dim=128):
        super(Net, self).__init__()
        self.encoder = []
        for name, module in resnet50(pretrained=True).named_children():
            if name == 'conv1':
                module = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
            if not isinstance(module, nn.Linear) and not isinstance(module, nn.MaxPool2d):
                self.encoder.append(module)
        # encoder
        self.encoder = nn.Sequential(*self.encoder)
        if head == 'linear':
            self.head = nn.Linear(dim_in, feat_dim)

        elif head == 'mlp':
            self.head = nn.Sequential(nn.Linear(2048, 512, bias=False), nn.BatchNorm1d(512),
                               nn.ReLU(inplace=True), nn.Linear(512, feat_dim, bias=True))
         
        else:
            raise NotImplementedError(
                'head not supported: {}'.format(head))

    def forward(self, x):
        feat = self.encoder(x)
        feat = torch.flatten(feat, start_dim=1)
        feat = F.normalize(self.head(feat), dim=-1)
        return feat


I am using mlp head , so the output is of dimension [batchsize,128] , as feat_dim=128 .