In " MYNETWORK " , I used a softmax function which I realized now works only for 2 values , what function should I replace it with .

import torch
import torch.nn as nn
import torch.nn.init as init


class mynetwork(nn.Module):

    def __init__(self, nfeats, width):
        super(mynetwork, self).__init__()
        
        self.width = width 
        self.nfeats = nfeats 
        
        mlp_modules = []
        for k in range(len(width)):
            mlp_modules.append(nn.Linear(width[k-1] if k>0 else nfeats, width[k]))
            mlp_modules.append(nn.ReLU(True))
        mlp_modules.append(nn.Linear(width[-1], 10))
        self.mlp = nn.Sequential(*mlp_modules)

    def forward(self, data):
        out = self.mlp(data).squeeze()        
        return nn.functional.softmax(out)