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)