I want to write a train function in a class for training a model; The following code reported an error; can anyone give me a hint for solving this issue?

import numpy as np
import os
import sys
sys.executable
sys.version
##define a neuralnet class
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

os.chdir("/Users/zhangzhongheng/Downloads/")
os.getcwd()

#Download training data from open datasets.

training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
)

Download test data from open datasets.

test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor(),
)

batch_size = 64

Create data loaders.

train_dataloader = DataLoader(training_data, batch_size=batch_size)
test_dataloader = DataLoader(test_data, batch_size=batch_size)

for X, y in test_dataloader:
    print(f"Shape of X [N, C, H, W]: {X.shape}")
    print(f"Shape of y: {y.shape} {y.dtype}")
    break

Get cpu or gpu device for training.

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")

Define model

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10)
        )
    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits
    def model_train(self,dataloader):
        self.train()
        size = len(dataloader.dataset)
        optimizer = torch.optim.SGD(self.parameters(), lr=1e-3)
        for batch, (X, y) in enumerate(dataloader):
            X, y = X.to(device), y.to(device)
        # Compute prediction error
            pred = self.forward(X)
            loss = nn.CrossEntropyLoss(pred, y)
        # Backpropagation
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if batch % 100 == 0:
                loss, current = loss.item(), batch * len(X)
                print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")
Model = NeuralNetwork()
epochs = 5
for t in range(epochs):
    print(f"Epoch {t+1}n-------------------------------")
    Model.model_train(train_dataloader)
print("Done!")

The above code reported the following error:

Epoch 1
-------------------------------
RuntimeError: Boolean value of Tensor with more than one value is ambiguous
Done!