You could create a model with two output neurons (e.g. via nn.Linear) and setup a multi-label classification use case using nn.BCEWithLogitsLoss.
Since the model output would be logits, you could apply torch.sigmoid on them to get the probabilities for each class.
Here is a very simple example:
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = nn.Linear(128, 64)
        self.fc2 = nn.Linear(64, 40)
        
    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = MyModel()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.BCEWithLogitsLoss()

data = torch.randn(64, 128)
target = torch.randint(0, 2, (64, 40)).float()

for epoch in range(100):
    optimizer.zero_grad()
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    print('epoch {}, loss {}'.format(epoch, loss.item()))

# check predictions
output = model(data)
probs = torch.sigmoid(output)
print(probs)
print(target)