I have chnaged the keras model and pytorch model and also incorporated the chnages you have suggested
#create a keras model
model=Sequential()

model.add(Dense(units=256,activation='relu',input_shape=[21])) 

model.add(Dense(units=128,activation='relu'))

model.add(Dense(units=64,activation='relu')) 

model.add(Dense(units=32,activation='relu'))

model.add(Dense(units=10,activation='softmax')) 

Here is keras summary
keras948×559 17.8 KB
Here is pytorch model
#we have create the same architecure in pytorch as we did in keras
class nonseqV2(torch.nn.Module):

    def __init__(self):

        super(nonseqV2,self).__init__()

        self.fcLayer1=torch.nn.Linear(in_features=21, out_features=256) 

        self.fcLayer2=torch.nn.Linear(in_features=256, out_features=128)

        self.fcLayer3=torch.nn.Linear(in_features=128, out_features=64)

        self.fcLayer4=torch.nn.Linear(in_features=64, out_features=32)

        self.fcLayer5=torch.nn.Linear(in_features=32, out_features=10)

        self.relu=torch.nn.ReLU()

        self.softmax=torch.nn.Softmax(dim=1)

        

    def forward(self,x):

        x=self.fcLayer1(x)

        x=self.relu(x)

        x=self.fcLayer2(x)

        x=self.relu(x)

        x=self.fcLayer3(x)

        x=self.relu(x)

        x=self.fcLayer4(x)

        x=self.relu(x)

        x=self.fcLayer5(x)

        #x=self.softmax(x)

        return x

Pytorch summary is in the attached image
pytorch993×655 21 KB
The loss obtained using Keras is 1.4 and in PyTorch its 4.0.
I am expecting Pytorch output to be closer to keras output