#!/usr/bin/env python3
from transformers import BertModel
import torch
import time


model = BertModel.from_pretrained("bert-base-uncased")
model.to("cuda")

input_ids = torch.ones((16, 256), dtype=torch.long)
input_ids = input_ids.to("cuda")

model.requires_grad_(False)

start_time = time.time()

for _ in range(5):
    with torch.no_grad():
        logits = model(input_ids)

print(time.time() - start_time)