And I used resnet18 to test. amp is faster, the memory usage is also lower.
def train_resnet18():
    x = torch.randn(64, 3, 224, 224, device='cuda')
    y = torch.randn(64, 1000, device='cuda')
    default_model = resnet18(pretrained=False)
    default_model.cuda()
    default_model.train()
    opt = torch.optim.SGD(default_model.parameters(), lr=1e-3)
    loss_fn = torch.nn.MSELoss().cuda()
    start_timer()
    for _ in range(500):
        y_pred = default_model(x)
        loss = loss_fn(y, y_pred)

        loss.backward()
        opt.step()
        opt.zero_grad()
    end_timer_and_print('default resnet18 training')

    mixed_model = resnet18(pretrained=False)
    mixed_model.cuda()
    mixed_model.train()
    scaler = GradScaler()
    mix_opt = torch.optim.SGD(mixed_model.parameters(), lr=1e-3)
    loss_fn = torch.nn.MSELoss().cuda()
    start_timer()
    for _ in range(500):
        with autocast():
            pred = mixed_model(x)
            loss = loss_fn(y, pred)
        scaler.scale(loss).backward()
        scaler.step(mix_opt)
        scaler.update()
        mix_opt.zero_grad()
    end_timer_and_print('mixed resnet18 training')

default resnet18 training
Total execution time = 27.581 sec
Max memory used by tensors = 3741790208 bytes
mixed resnet18 training
Total execution time = 15.231 sec
Max memory used by tensors = 3247211520 bytes