Is there a way that I can have my model output a specific shape? Say if my input (ground truth) has a shape of [15, 3], however, the output (x_kp) size from the model is [15, 1].
class Model(nn.Module):
    def __init__(self, num_classes, batch_size):
        super(Model, self).__init__()
        self.conv1 = nn.Conv3d(in_channels=256, out_channels=512, kernel_size=(3, 3, 3), stride=(1, 1, 1))
        self.conv2 = nn.Conv3d(in_channels=512, out_channels=512, kernel_size=(3, 3, 3), stride=(1, 1, 1))
        self.conv3 = nn.Conv3d(in_channels=512, out_channels=512, kernel_size=(3, 3, 3), stride=(1, 1, 1))
        self.conv4 = nn.Conv3d(in_channels=512, out_channels=512, kernel_size=(3, 3, 3), stride=(1, 1, 1))
        self.conv5 = nn.Conv3d(in_channels=512, out_channels=512, kernel_size=(3, 3, 3), stride=(1, 1, 1))
        self.conv6 = nn.Conv3d(in_channels=512, out_channels=512, kernel_size=(3, 3, 3), stride=(1, 1, 1))
        self.conv7 = nn.Conv3d(in_channels=512, out_channels=512, kernel_size=(3, 3, 3), stride=(1, 1, 1))
        self.conv8 = nn.Conv3d(in_channels=512, out_channels=512, kernel_size=(3, 3, 3), stride=(1, 1, 1))
        self.cls_score = nn.Linear(471040, num_classes)
        self.kp_score = nn.Linear(471040, 17*batch_size)
        
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = F.relu(self.conv4(x))
        x = F.relu(self.conv5(x))
        x = F.relu(self.conv6(x))
        x = F.relu(self.conv7(x))
#         x = F.relu(self.conv8(x))
        x = x.flatten(start_dim=1)
        x_label = F.relu(self.cls_score(x))
        x_kp = F.relu(self.kp_score(x))
        return x_label, x_kp
    
model = Model(num_classes=2, batch_size=batch_size).to(device)

My criterion and loss is calculated using the following:
pred_label, pred_kp = model(cat_features)

kp_criterion = nn.SmoothL1Loss()
label_criterion = nn.CrossEntropyLoss()

kp_loss = kp_criterion(torch.mean(pred_kp), torch.mean(cat_kp).to(device))
label_loss = label_criterion(pred_label, crossing_labels)