TorchVista

import torch
from torchvista import trace_model
from torchvision import models

model = models.vit_b_16(weights=None)
example_input = torch.randn(1, 3, 224, 224)
trace_model(model, example_input, forced_module_tracing_depth=5)