Append values from the iterator into objects

import torch
import torchvision
from torch.utils.tensorboard import SummaryWriter
from torchvision import datasets, transforms

# Writer will output to ./runs/ directory by default
writer = SummaryWriter()

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
trainset = datasets.MNIST('mnist_train', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
model = torchvision.models.resnet50(True)
# Have ResNet model take in grayscale rather than RGB
model.conv1 = torch.nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)

device=torch.device('mps' if torch.has_mps else 'cpu')
for i in ["mps","cpu"]:
    model=model.to(device)
    images, labels = next(iter(trainloader))
    images=images.to(device)

    grid = torchvision.utils.make_grid(images.to(device))
    writer.add_image('images', grid, 0)
    writer.add_graph(model.to(device), images.to(device))
    writer.close()

I have this code where I would like to append the values of the iterator (mps & cpu) in the writer.add_image() & writer.add_graph() object, so that I can differentiate which originated from which run. In other words, I would like to see images_mps & images_cpu in the [tag:tensorboard] dashboard (by some modification in the writer.add_image('images', grid, 0) line?).

Sorry if it's noob question, but appreciate any help/suggestion!

Try to make 2 SummaryWriter specifying 2 log_dir, then use one writer for mps and the other for `cpu
https://pytorch.org/docs/stable/tensorboard.html#torch.utils.tensorboard.writer.SummaryWriter