Fine tuned efficientnet doesn't train

Hello! Can someone explain me why this model can’t train?

class EfficientNetClassifier(pl.LightningModule):
    def __init__(self, num_classes, transfer=True, freeze="last"):
        super().__init__()

        self.efficientNet = models.efficientnet_b4(pretrained=transfer)
        linear_size = 1792
        self.efficientNet.classifier[1] = nn.Sequential(
            nn.Linear(linear_size, 300),
            nn.BatchNorm1d(300),
            nn.LeakyReLU(),
            nn.Linear(300, 50)
        )
        
class LightningClassifier(pl.LightningModule):
    def __init__(self, lr_rate=1e-5, freeze="last"):
        super().__init__()

        self.model = EfficientNetClassifier(50, True, freeze)

        self.lr_rate = lr_rate

    def forward(self, x):
        return self.model(x)

    def cross_entropy_loss(self, logits, labels):

        return F.nll_loss(logits, labels)

    def training_step(self, train_batch, batch_idx):
        x, y = train_batch
        logits = self.forward(x)
        loss = self.cross_entropy_loss(logits, y)

        logs = {"train_loss": loss}

        acc = torch.sum(logits.argmax(axis=1) == y) / y.shape[0]
        self.log("train_acc", acc, on_step=True, prog_bar=True)
        self.log("train_loss", loss, on_step=True, prog_bar=True)

        return {"loss": loss, "log": logs}

    def validation_step(self, val_batch, batch_idx):
        x, y = val_batch
        logits = self.forward(x)
        loss = self.cross_entropy_loss(logits, y)
        acc = torch.sum(logits.argmax(axis=1) == y) / y.shape[0]

        self.log("val_loss", loss, on_step=True, on_epoch=True, prog_bar=True)
        self.log("val_acc", acc, on_step=True, on_epoch=True, prog_bar=True)

        return {"val_loss": loss, "val_acc": acc}

I chcked everything in code few times but didn’t find mistake. I thought that is because of dropout in efficientnet model, but as I know pytorch lightning include model.train() and model.eval() in training_step(), validation_step() and turns off dropout on evaluation