This code is working fine for me:
temp = nn.Sequential(
          nn.Dropout(p=0.5),
          nn.Linear(in_features=512, out_features=128),
          nn.ReLU(),
          nn.Linear(in_features=128, out_features=17, bias=True),
        )

classifier = models.resnet34(pretrained=True)
classifier.fc = temp

device = 'cuda'  
classifier.to(device)
optimizer = torch.optim.SGD(classifier.parameters(), lr=1e-3)
scaler = torch.cuda.amp.GradScaler()

data = torch.randn(2, 3, 224, 224, device=device)
target = torch.randint(0, 2, (2, 17), device=device).float()
criterion = nn.BCEWithLogitsLoss()

for epoch in range(10):
    optimizer.zero_grad()
    with torch.cuda.amp.autocast():
        logits = classifier(data)
        loss = criterion(logits, target)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    # Updates the scale for next iteration.
    scaler.update()
    print('epoch {}, loss {:.3f}'.format(epoch, loss.item()))