My training loop look like code below in actual, the code above was just an example, apologies for the confusion:
def train_classifier(classifier, train_loader, optimizer, criterion):
    classifier.half()
    classifier.train()
    loss = 0.0
    losses = []
    for i, (images, labels) in enumerate(train_loader):
        classifier.half()
        images, labels = images.to(device), labels.float().to(device)
        images = images.half()
        optimizer.zero_grad()
        logits = classifier(images)
        logits = logits.float()
        loss = criterion(logits, labels)
        loss = loss.float()
        loss.backward()
        classifier.float()
        optimizer.step()
        losses.append(loss)
    return torch.stack(losses).mean().item()

Any idea what could be wrong?