Dimension out of range (expected to be in range of [-1, 0], but got 1)

Question

I have been trying to implement MNIST dataset on a simple Neural Network(784,512,128,10) using cross-entropy loss. I am using Keras for getting the MNIST dataset. But I'm getting an error :

RuntimeError: 1D target tensor expected, multi-target not supported

When I have main model as:

for epoch in range(num_epochs):
  for x,y in train_data:
    x=Variable(x)
    y=Variable(y)
    print(x.shape)
    y_pred=model(x)
    optimizer.zero_grad()
    loss=criterion(y_pred,y)
    loss.backward()
    optimizer.step()

So, to remove that error I implemented that:

y=y[0][0:]
y_pred=y_pred[0][0:]
loss=criterion(y_pred,y)

But after that I'm getting this error:

IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)

I read many articles regarding how to solve this error but none helped.

Is this error coming because of Keras dataset? or Is something wrong in my code can someone please help find the error? My Code:

import torch
import torch.nn as nn
import numpy as np
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader,Dataset
import keras
import torch.nn.functional as F
from torch.autograd import Variable

class Netz(nn.Module):
  def __init__(self,n_input_features):
    super(Netz,self).__init__()
    self.linear=nn.Linear(784,512,bias=True)
    self.l1=nn.Linear(512,128,bias=True)
    self.l2=nn.Linear(128,10,bias=True)
    self.relu=nn.ReLU()
    self.relu2=nn.ReLU()
    self.softmax=nn.Softmax(dim=-1)
  def forward(self,x):
    # x=x.view(-1,784)
    x=self.relu(self.linear(x))
    x=self.relu2(self.l1(x))
    x=self.softmax(self.l2(x))
    return x

model=Netz(784)

class Data(Dataset):
    def __init__(self):
        self.x=x_train
        self.y=y_train
        self.len=self.x.shape[0]
    def __getitem__(self,index):
        return self.x[index],self.y[index]

mnist = keras.datasets.mnist
#Copying data
(x_train, y_train),(x_test, y_test) = mnist.load_data()
#One-hot encoding the labels
y_train = keras.utils.to_categorical(y_train,10)
y_test = keras.utils.to_categorical(y_test,10)
#Flattening the images
x_train_reshaped = x_train.reshape((60000,784))
x_test_reshaped = x_test.reshape((10000,784))
#Normalizing the inputs
x_train = x_train_reshaped/255.0 
x_test = x_test_reshaped/255.0

x_train=torch.from_numpy(x_train.astype(np.float32))
x_test=torch.from_numpy(x_test.astype(np.float32))
y_train=torch.from_numpy(y_train.astype(np.float32))
y_test=torch.from_numpy(y_test.astype(np.float32))

criterion=nn.CrossEntropyLoss()
print(criterion)
optimizer=torch.optim.SGD(model.parameters(),lr=0.05)
dataset=Data()
train_data=DataLoader(dataset=dataset,batch_size=1,shuffle=False)

num_epochs=5
for epoch in range(num_epochs):
  for x,y in train_data:
    x=Variable(x)
    y=Variable(y)
    y_pred=model(x)
    optimizer.zero_grad()
    y=y[0][0:]
    y_pred=y_pred[0][0:]
    loss=criterion(y_pred,y)
    loss.backward()
    optimizer.step()