Manual test loop results are awful compared to trainer.test() results
🐛 Bug
Using MNIST as an example, when I train using EfficientNet (I've also used other architectures from pretrainedmodels and I have the same issue) I get ~77% Accuracy after two epochs on validation and when I use trainer.test(). However, if I use my own testing loop my Accuracy is ~10%... as good as random. If I extract a batch from the test_dataloader and pass it through the model i.e. out = model(x), it predicts the same value for nearly every input. It's almost as if model(x) is using the weights before training started instead of the weights after training ends. Am I doing something wrong when trying to get predictions via model(x) or is this a bug?
To Reproduce
See code sample
Code sample
import os
import sys
import gc

import cv2
import numpy as np
import pandas as pd
import pickle as pkl
from argparse import ArgumentParser
import argparse
from tqdm import tqdm
from PIL import Image
Image.MAX_IMAGE_PIXELS = 1000000000000

from torch import nn
from torch.utils.data import Dataset, DataLoader, RandomSampler
import torch.nn.functional as F
from torchvision import transforms
from torchvision.transforms import functional as TF
from torchvision.datasets import MNIST
import torch

import pytorch_lightning as pl
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import EarlyStopping
from pytorch_lightning.logging.test_tube import TestTubeLogger
from pytorch_lightning.logging import TensorBoardLogger

from efficientnet_pytorch import EfficientNet

IMG_HEIGHT = 300
IMG_WIDTH = 300

def parse_args(args):

    parser = argparse.ArgumentParser()

    parser.add_argument("--backbone",
                        help="Backbone to use",
                        type=str,
                        default="efficientnet-b4")
    parser.add_argument("--backbone_weights",
                        type=str,
                        default="imagenet")
    parser.add_argument("--batch_size",
                        help="Batch size",
                        type=int,
                        default=12)
    parser.add_argument("--acc_grad",
                        help="Accumulate gradient for acc_grad epochs",
                        type=int,
                        default=4)
    parser.add_argument("--gpus",
                        help="GPUs to use. Either an int or comma-separated list e.g. '0,1,2'",
                        type=int,
                        default=1)
    parser.add_argument("--num_workers",
                        help="Number of workers to use in DataLoader.",
                        type=int,
                        default=8)    

    args = parser.parse_args(args)
    return(args)


def accuracy(output, target, topk=(1,)):
    """Computes the accuracy over the k top predictions for the specified values of k"""
    with torch.no_grad():
        maxk = max(topk)
        batch_size = target.size(0)

        _, pred = output.topk(maxk, 1, True, True)
        pred = pred.t()
        correct = pred.eq(target.view(1, -1).expand_as(pred))

        res = []
        for k in topk:
            correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
            res.append(correct_k.mul_(100.0 / batch_size))
        return res


class TrainSystem(pl.LightningModule):

    def __init__(self, hparams, test_data=None):
        super(TrainSystem, self).__init__()

        self.hparams = hparams

        self.model = EfficientNet.from_pretrained(self.hparams.backbone, advprop=False, num_classes=self.hparams.n_classes)
        self.model = self.model.float()

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

    def training_step(self, batch, batch_idx):
        x, y = batch
        out = self.forward(x)
        if isinstance(out, dict):
            out = out['out']
        celoss = nn.CrossEntropyLoss()
        loss = celoss(out, y)
        tensorboard_logs = {'train_loss': loss}
        return {'loss': loss, 'log': tensorboard_logs}

    def validation_step(self, batch, batch_idx):
        x, y = batch
        out = self.forward(x)
        if isinstance(out, dict):
            out = out['out']
        celoss = nn.CrossEntropyLoss()
        loss = celoss(out, y)
        out_max = torch.argmax(out, dim=1)
#         if batch_idx % 10 == 0:
#             print('out: ', out)
#             print("out_max: ", out_max)
#             print("y: ", y)
        acc = accuracy(out, y)[0]
        return {'val_loss': loss,
               'val_acc': acc}

    def validation_end(self, outputs):
        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
        avg_acc = torch.stack([x['val_acc'] for x in outputs]).mean()
        print(avg_acc)
        tensorboard_logs = {'val_loss': avg_loss,
                           'val_acc': avg_acc}
        return {'avg_val_loss': avg_loss, 
                'avg_val_acc': avg_acc,
                'log': tensorboard_logs}
            
    def test_step(self, batch, batch_idx):
        # OPTIONAL
        x, y = batch
        out = self.forward(x)
        if isinstance(out, dict):
            out = out['out']
        celoss = nn.CrossEntropyLoss()
        loss = celoss(out, y)
        out_max = torch.argmax(out, dim=1)
#         if batch_idx % 10 == 0:
#             print('out: ', out)
#             print("out_max: ", out_max)
#             print("y: ", y)
        acc = accuracy(out, y)[0]
        return {'test_loss': loss,
               'test_acc': acc}

    def test_end(self, outputs):
        # OPTIONAL
        avg_loss = torch.stack([x['test_loss'] for x in outputs]).mean()
        avg_acc = torch.stack([x['test_acc'] for x in outputs]).mean()
        print(avg_acc)
        tensorboard_logs = {'test_loss': avg_loss,
                           'test_acc': avg_acc}
        return {'avg_test_loss': avg_loss, 
                'avg_test_acc': avg_acc,
                'log': tensorboard_logs}

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.hparams.lr)

    def train_dataloader(self):
        # REQUIRED
        return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.Compose([transforms.Resize((IMG_WIDTH, IMG_HEIGHT)), transforms.Grayscale(3), transforms.ToTensor()])), batch_size=self.hparams.batch_size, num_workers=self.hparams.num_workers, shuffle=True, pin_memory=True, drop_last=True)

    def val_dataloader(self):
        # OPTIONAL
        return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.Compose([transforms.Resize((IMG_WIDTH, IMG_HEIGHT)),transforms.Grayscale(3), transforms.ToTensor()])), batch_size=self.hparams.batch_size, num_workers=self.hparams.num_workers, shuffle=False, pin_memory=True, drop_last=True)

    def test_dataloader(self):
        # OPTIONAL
        return DataLoader(MNIST(os.getcwd(), train=False, download=True, transform=transforms.Compose([transforms.Resize((IMG_WIDTH, IMG_HEIGHT)),transforms.Grayscale(3), transforms.ToTensor()])), batch_size=self.hparams.batch_size, num_workers=self.hparams.num_workers, shuffle=False, pin_memory=True, drop_last=False)



args = ['--backbone', 'efficientnet-b4','--batch_size','48','--acc_grad','1','--gpus','4']
args = parse_args(args)
print(args)
    
out_dir = f"./saved/chk/classifier_{args.backbone}"
    
hparams = argparse.Namespace(**{'backbone': args.backbone,
                                    'encoder_weights': args.backbone_weights,
                                    'lr': 0.0001,
                                    'resize': (IMG_WIDTH, IMG_HEIGHT),
                                    'n_classes': 10,
                                    'batch_size': args.batch_size,
                                    'num_workers': args.num_workers})

logger = TestTubeLogger(
                        save_dir=out_dir,
                        version=f'MNIST',
                        name='lightning_logs'
                    )

early_stop_callback = EarlyStopping(
            monitor='val_acc',
            patience=20,
            # strict=False,
            verbose=True,
            mode='max'
        )

trainer = Trainer(logger=logger,
                          default_save_path=out_dir,
                          accumulate_grad_batches=args.acc_grad,
                          early_stop_callback=early_stop_callback,
                          use_amp=True,
                          amp_level='O1',
                          gpus=args.gpus,
                          max_epochs=2,
                          distributed_backend='dp') 

model = TrainSystem(hparams)

trainer.fit(model)
        
trainer.test()
# tensor(77.1830, device='cuda:0')
# 77% Accuracy after 2 epochs! Not great, but clearly it's learning something

def get_predictions(model, dataloader):
        model.freeze()
        #model = model.cuda()
        preds = []
        accs = []
        with torch.no_grad():
            for x, y in dataloader:
                x = x.cuda()
                y = y.cuda()
                #print(y)
                out0 = model(x)
                acc = accuracy(out0, y)
                out0 = torch.topk(out0, k=1).indices.tolist()
                #print(out0)
                preds.extend(out0)
                accs.extend(acc)

        preds = np.asarray(preds)
        accur = np.mean([a.cpu().numpy()[0] for a in accs])
        return preds, accur

test_dl = trainer.test_dataloaders[0]
pred, accur = get_predictions(model, test_dl)
accur
# 9.808613%
# Now only 9.8% Accuracy -- As good as random guessing

If I predict on a single batch from the test_dataloader almost all of the predictions are of the same class which would explain why Accuracy is ~10%.
x, y = next(test_dl.__iter__())
out = model(x.cuda())
torch.argmax(out, dim=1)
# tensor([7, 5, 5, 5, 5, 5, 7, 5, 5, 5, 5, 5, 7, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5], device='cuda:0')
# It predicts 5 nearly all the time

Expected behavior
I would expect much better results from model(x) and my prediction loop.
Environment
cuda:
GPU:
Tesla P100-PCIE-16GB
available:           True
version:             10.1
packages:
numpy:               1.18.2
pyTorch_debug:       False
pyTorch_version:     1.4.0
pytorch-lightning:   0.7.1
tensorboard:         2.1.1
tqdm:                4.42.0
system:
OS:                  Linux
architecture:
64bit
processor:           
python:              3.6.6
version:             #1 SMP Tue Feb 4 23:02:59 UTC 2020