写一个能够在pycharm里面动态展示每个epoch的训练损失的pytorch代码

这里是一个能够在pycharm里面动态展示每个epoch的训练损失的pytorch代码：

import torch
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt

# 定义超参数
batch_size = 64
learning_rate = 1e-3
epochs = 10

# 加载数据集
train_dataset = MNIST(root='./data', train=True, transform=ToTensor(), download=True)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

# 定义模型
model = torch.nn.Sequential(
    torch.nn.Linear(784, 512),
    torch.nn.ReLU(),
    torch.nn.Linear(512, 256),
    torch.nn.ReLU(),
    torch.nn.Linear(256, 10)
)

# 定义损失函数和优化器
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# 训练模型
train_losses = []
for epoch in range(epochs):
    epoch_loss = 0
    for batch_idx, (data, target) in enumerate(train_loader):
        # 将输入数据展开成一维向量
        data = data.view(-1, 784)

        # 前向传播
        output = model(data)

        # 计算损失
        loss = criterion(output, target)

        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # 累加损失
        epoch_loss += loss.item()

    # 计算每个epoch的平均损失
    epoch_loss /= len(train_loader)

    # 打印每个epoch的损失
    print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, epochs, epoch_loss))

    # 将每个epoch的损失保存到列表中
    train_losses.append(epoch_loss)

    # 动态展示每个epoch的训练损失
    plt.plot(train_losses)
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.title('Training Loss')
    plt.pause(0.1)

这段代码使用了Matplotlib库来动态展示每个epoch的训练损失。在每个epoch结束后，程序会将当前的训练损失保存到列表中，并使用Matplotlib绘制出每个epoch的训练损失曲线。由于在绘制完成后程序会暂停0.1秒，因此可以看到每个epoch的训练损失曲线是逐渐更新的