First diverging operator message for vit_b_16 and torch.utils.tensorboard

Question

I can't get torch.utils.tensorboard to write_graph for vit_b_16 model.

Here is example code:

import torch
from torchvision.models import get_model
from torch.utils import tensorboard

# create example inputs
fake_images = torch.ones((32, 3, 224, 224))

vit_model = get_model("vit_b_16", num_classes=1000, image_size=224, weights="IMAGENET1K_V1")

writer = tensorboard.SummaryWriter("tb_writer_outputs")
writer.add_graph(vit_model, fake_images)

# the seemingly problematic messages about "First diverging operator" occur
# when writer.add_graph is ran, so I don't think the follow 2 lines are an issue
writer.flush()
writer.close()

There is a lot of text which prints to the terminal. Here is the end of the messages: message which says "First diverging operator" followed by some other stuff

torch.__version__ is 2.0.0+cu117

torchvision.__version__ is 0.15.1+cu117

What does the above example code not work to write the graph?

Can tensorboard.write_graph work with vit models like vit_b_16 in torchvision.models?