and for backward pass:
from torch.utils import benchmark
import torch
from torch.cuda.amp import autocast
from transformers import BertForSequenceClassification, BertTokenizer

model = BertForSequenceClassification.from_pretrained("bert-base-uncased").to(0)

tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
inputs = tokenizer("Hello, my dog is cute", return_tensors="pt").to(0)
labels = torch.tensor([1]).unsqueeze(0).to(0)

def autocast_backward(model, inputs, labels, fast_dtype=None):
    if fast_dtype is None:
        outputs = model(**inputs, labels=labels)
        loss = outputs.loss
        loss.backward()
        return

    with autocast(fast_dtype=fast_dtype):
        outputs = model(**inputs, labels=labels)
        loss = outputs.loss
        loss.backward()

    return

t0 = benchmark.Timer(
    stmt="autocast_backward(model, inputs, labels)",
    setup="from __main__ import autocast_backward",
    globals={"inputs": inputs, "model": model, "labels": labels,},
)

t1 = benchmark.Timer(
    stmt="autocast_backward(model, inputs, labels, torch.float16)",
    setup="from __main__ import autocast_backward",
    globals={"inputs": inputs, "model": model, "labels": labels,},
)
t2 = benchmark.Timer(
    stmt="autocast_backward(model, inputs, labels, torch.bfloat16)",
    setup="from __main__ import autocast_backward",
    globals={"inputs": inputs, "model": model, "labels": labels,},
)

print(t0.timeit(100))  # prints 20.52
print(t1.timeit(100))  # prints 22.33
print(t2.timeit(100))  # prints 26.43 :(