pip install fasterbench
from torchvision.models import resnet18 model = resnet18() dummy_input = torch.randn(64, 3, 224, 224) benchmark(model, dummy_input)