You would have to adapt the model as described in my last post:



 ptrblck:

You could use different heads in your model and the corresponding loss functions for the different targets.
I.e. your classification head could be a linear layer which outputs the class logits, while the bboxes would be provided by another linear layer on top of the model. This would also mean that your model would return more than one output tensor.


Here is a small dummy example:
class MyModel(nn.Module):
    def __init__(self, nb_classes):
        super(MyModel, self).__init__()
        self.base = nn.Sequential(
            nn.Conv2d(3, 6, 3, 1, 1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(6, 12, 3, 1, 1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        
        self.out_labels = nn.Linear(12*6*6, nb_classes)
        self.out_bbox = nn.Linear(12*6*6, 4)
        
    def forward(self, x):
        x = self.base(x)
        x = x.view(x.size(0), -1)
        x_labels = self.out_labels(x)
        x_bbox = self.out_bbox(x)
        return x_labels, x_bbox

model = MyModel(nb_classes=10)
x = torch.randn(1, 3, 24, 24)
output_labels, output_bbox = model(x)