Yes, you are right. Following your suggestion I have done the following, please see if I am doing that right:
1- I created a module list of linear layers equal to the number of classes, i.e. 10 linear layers
2- Each linear layer linear( feature_dim, 1)
3- I pass the features to all the modules, and calculate the loss with respect to each independently.
the model:
class one_vs_all(nn.Module):
    def __init__(self, input_dim=32,out_dim=1, num_classes=10, dropout=0):
        super(one_vs_all, self).__init__()
        # not the best model...
        self.input_dim = input_dim
        self.out_dim =out_dim
        self.sigmoid = nn.Sigmoid()
        self.num_classes=num_classes
        self.base_model= nn.Linear(self.input_dim,self.out_dim) 
        self.classifiers =  nn.ModuleList([self.base_model for i in range(self.self.num_classes)])
    def forward(self,class_index,x):
        predictions = self.classifiers[class_index](x)
        logits =self.sigmoid(predictions)
        return logits

and here is the training function
def train(model, train_dl, optimizer, criterion,config,device):
    model.train()
    epoch_loss = 0
    class_accuracy = {f'class{k+1}': [] for k in range(len(num_classes))}
    features =[];train_labels=[];total_pred= []
    Loss=0
    for inputs, labels in train_dl:
        src = inputs.to(device)
        labels = labels.to(device)
        #loss and score
        # one_vs_all
        for idx in range(num_classes):
            optimizer.zero_grad()
            tgt_class_label = torch.eq(labels, idx).float()
            logits        = model(idx,inputs)
            class_loss = binary_cross_entropy(logits,tgt_class_label)
            Loss+=class_loss
            class_accuracy[f'D{idx+1}'].append() = accuracy_score(tgt_class_label.cpu().detach().numpy(),pred.argmax(axis=1).cpu().detach().numpy())
            class_loss.backward()
            optimizer.step()
        epoch_loss += Loss.item()
        epoch_accuracy += Accuracy
        total_pred.append(pred.argmax(axis=1))
        features.append(feas)
        train_labels.append(labels)
    return epoch_loss / len(train_dl), epoch_accuracy/len(train_dl), torch.cat(features),torch.cat(total_pred),torch.cat(train_labels)```