Thank you very much for your help!
Now I try to add localization. After looking for some information on the internet, this is the code:
class CNN(nn.Module):
    def __init__(self, num_classes):
        super(CNN, self).__init__()
        self.model_resnet = models.resnet18(pretrained=True)
        fc_inputs = self.model_resnet.fc.in_features
        self.model_resnet.fc = nn.Identity()

        self.labels_model = nn.Sequential(
            nn.Linear(fc_inputs, 256),
            nn.ReLU(),
            nn.Dropout(0.4),
            nn.Linear(256, num_classes), 
            nn.LogSoftmax(dim=1) # For using NLLLoss()
        )
  
        self.conv1 = nn.Conv2d(fc_inputs, 256, kernel_size=3, padding=1)
        self.act1 = nn.ReLU()
        self.conv2 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        self.act2 = nn.ReLU()
        self.conv3 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        self.act3 = nn.ReLU()
        self.conv4 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        self.act4 = nn.ReLU()
        self.output = nn.Conv2d(256, 4, kernel_size=3, padding=1)


    def forward(self, images):
        with torch.no_grad():
            features = self.model_resnet(images)

        labels = self.labels_model(features)      
      
        box = self.conv1(features)
        box = self.act1(box)
        box = self.conv2(box)
        box = self.act2(box)
        box = self.conv3(box)
        box = self.act3(box)
        box = self.conv4(box)
        box = self.act4(box)
        box = self.output(box)
        box = out.permute(0, 2, 3, 1)
        box = out.contiguous().view(out.shape[0], -1, 4)

        return labels, box

But I get the next error:
“RuntimeError: Expected 4-dimensional input for 4-dimensional weight 256 512, but got 2-dimensional input of size [32, 512] instead”
Thank you again!