# 创建TensorBoard写入器 writer = SummaryWriter()

from torch.utils.tensorboard import SummaryWriter

# 创建TensorBoard写入器
writer = SummaryWriter()

# 训练模型并记录损失值
for epoch in range(1000):
    optimizer.zero_grad()
    output = model(X)
    loss = criterion(output, y)
    loss.backward()
    optimizer.step()

    writer.add_scalar('Loss/train', loss.item(), epoch)

# 关闭写入器
writer.close()