BERT encoding layer produces same output for all inputs during evaluation (PyTorch)
I don't understand why my BERT model returns the same output during evaluation. The output of my model during training seems correct, as the values were different, but is totally the same during evaluation.


Here is my BERT model class 

class BERTBaseUncased(nn.Module):
    def __init__(self):
        super(BERTBaseUncased, self).__init__()
        self.bert = BertModel.from_pretrained("bert-base-uncased")
        self.bert_drop = nn.Dropout(0.3)
        self.out = nn.Linear(768, 4)

    def forward(self, ids, mask, token_type_ids):
        _, o2 = self.bert(ids, attention_mask=mask, token_type_ids=token_type_ids) # Use one of the outputs
        bo = self.bert_drop(o2)
        return self.out(bo)


My dataset class

class BERTDataset:
    def __init__(self, review, target, tokenizer, classes=4):
        self.review = review
        self.target = target
        self.tokenizer = tokenizer
        self.max_len = max_len
        self.classes = classes

    def __len__(self):
        return len(self.review)

    def __getitem__(self, item):
        review = str(self.review)
        review = " ".join(review.split())

        inputs = self.tokenizer.encode_plus(review, None, add_special_tokens=True, max_length= self.max_len,
                                            pad_to_max_length=True, return_token_type_ids=True,
                                            return_attention_masks=True)

        ids = inputs["input_ids"]
        mask = inputs["attention_mask"]
        token_type_ids = inputs["token_type_ids"]

        return {
            'ids': torch.tensor(ids, dtype=torch.long),
            'mask': torch.tensor(mask, dtype=torch.long),
            'token_type_ids': torch.tensor(token_type_ids, dtype=torch.long),
            'targets': torch.tensor(to_categorical(self.target[item], self.classes), dtype=torch.float)
        }


My evaluation function

def eval_fn(data_loader, model, device):
    model.eval()

    total_loss = 0.0

    with torch.no_grad():
        for bi, d in tqdm(enumerate(data_loader), total=len(data_loader)):
            ids = d['ids']
            token_type_ids = d['token_type_ids']
            mask = d['mask']
            targets = d['targets']

            ids = ids.to(device, dtype=torch.long)
            token_type_ids = token_type_ids.to(device, dtype=torch.long)
            mask = mask.to(device, dtype=torch.long)
            targets = targets.to(device, dtype=torch.float)

            outputs = model(
                ids=ids,
                mask=mask,
                token_type_ids=token_type_ids
            )

            loss = loss_fn(outputs, targets)
            total_loss += loss.item()


And my training function

def train_fn(data_loader, model, optimizer, device, scheduler):
    model.train()

    total_loss = 0.0

    for bi, d in tqdm(enumerate(data_loader), total=len(data_loader)):
        ids = d['ids']
        token_type_ids = d['token_type_ids']
        mask = d['mask']
        targets = d['targets']

        ids = ids.to(device, dtype=torch.long)
        token_type_ids = token_type_ids.to(device, dtype=torch.long)
        mask = mask.to(device, dtype=torch.long)
        targets = targets.to(device, dtype=torch.float)

        optimizer.zero_grad()

        outputs = model(
            ids=ids,
            mask=mask,
            token_type_ids=token_type_ids
        )

        loss = loss_fn(outputs, targets)
        total_loss += loss.item()
        loss.backward()

        optimizer.step()
        scheduler.step()

    return total_loss/len(data_loader)


Thanks!

In case anybody else has the problem, perhaps you forgot to use one of the recommended learning rates from the official paper:  5e-5, 3e-5, 2e-5
Gradients seem to polarize if the learning rate is too high, such as 0.01, causing repeatedly the same logits for the val set.
