Hi Rishav!
I haven’t looked at your code in detail, but I have two comments:



 Rishav_Sen:

class Net(nn.Module):
    def __init__(self, D_in, H1, H2, D_out):
        super(Net, self).__init__()
        self.linear1 = nn.Linear(D_in, H1)
        self.linear2 = nn.Linear(H1, H2)
        self.linear3 = nn.Linear(H2, D_out)
        self.softmax = nn.Softmax(dim = 0)

    def forward(self, x):
        x = (self.linear1(x))
        x = (self.linear2(x))
        x = (self.linear3(x))
        x = (self.softmax(x))
        return x



First, be aware that because you do not have any intervening
non-linearity between your Linear layers, your three Linears
collapse, in effect, into a single linear layer equivalent to
Linear (D_in, D_out).
You should place non-linear so-called “activations,” such as
ReLU or Tanh between the layers, e.g.:
    def forward(self, x):
        x = (self.linear1(x))
        x = torch.nn.functional.relu (x)
        x = (self.linear2(x))
        x = torch.nn.functional.relu (x)
        x = (self.linear3(x))
        return x



        self.softmax = nn.Softmax(dim = 0)
...
        x = (self.linear3(x))
        x = (self.softmax(x))
        return x
...
CrossELoss = nn.CrossEntropyLoss()
...
        z = model(x)
        loss = CrossELoss(z, y)



Second, your loss criterion, CrossEntropyLoss, has, in effect,
softmax() built in.  Therefore your model should not have a final
Softmax layer; just feed the output of your final Linear layer into
your loss function.
Best.
K. Frank