PL computes wrong accuracy with drop_last=False in PyTorch Geometric
🐛 Bug
PyTorch Lightning computes wrong accuracy when using a DataLoader with drop_last=False in PyTorch Geometric.
There seems to be an issue in which PL cannot determine the correct batch_size of mini-batches.
from typing import Optional

import torch
import torch.nn.functional as F
from torch.nn import Linear
from pytorch_lightning.metrics import Accuracy
from pytorch_lightning import (LightningDataModule, LightningModule, Trainer,
                               seed_everything)

from torch_geometric.data import DataLoader
from torch_geometric.datasets import TUDataset
from torch_geometric.nn import GCNConv, global_mean_pool


class Dataset(LightningDataModule):
    def __init__(self):
        super().__init__()

    def prepare_data(self):
        TUDataset('./data', name='MUTAG')

    def setup(self, stage: Optional[str] = None):
        dataset = TUDataset('./data', name='MUTAG')
        self.train_dataset = dataset[:3]
        self.val_dataset = dataset[3:6]
        self.test_dataset = dataset[6:9]

    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=2)

    def val_dataloader(self):
        return DataLoader(self.val_dataset, batch_size=2)

    def test_dataloader(self):
        return DataLoader(self.val_dataset, batch_size=2)


class GNN(LightningModule):
    def __init__(self):
        super().__init__()
        self.conv = GCNConv(7, 64)
        self.lin = Linear(64, 2)
        self.acc = Accuracy()

    def forward(self, x, edge_index, batch):
        x = self.conv(x, edge_index).relu()
        x = global_mean_pool(x, batch)
        return self.lin(x)

    def training_step(self, data, batch_idx):
        data = data.to(self.device)
        y_hat = self(data.x, data.edge_index, data.batch)
        train_loss = F.cross_entropy(y_hat, data.y)
        return train_loss

    def validation_step(self, data, batch_idx):
        data = data.to(self.device)
        y_hat = self(data.x, data.edge_index, data.batch)
        acc = self.acc(y_hat.softmax(dim=-1), data.y)
        self.log('val_acc', acc, on_step=False, on_epoch=True)
        return acc

    def test_step(self, data, batch_idx):
        data = data.to(self.device)
        y_hat = self(data.x, data.edge_index, data.batch)
        acc = self.acc(y_hat.softmax(dim=-1), data.y)
        print('batch_size', data.num_graphs, 'accuracy', acc, 'shape', y_hat.shape)
        self.log('test_acc', acc, on_step=False, on_epoch=True)
        return acc

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


def main():
    seed_everything(42)
    datamodule = Dataset()
    model = GNN()
    trainer = Trainer(max_epochs=1, progress_bar_refresh_rate=0)
    trainer.fit(model, datamodule=datamodule)
    trainer.test()


if __name__ == "__main__":
    main()
Here, I am using a dataset with 3 examples and utilize a batch_size of 2. In test_step, the accuracy of each individual mini-batch is:
batch_size 2 accuracy 0.5 shape [2, 2]
batch_size 1 accuracy 0.0 shape [1, 2]

while PyTorch Lightning reports an overall accuracy of 0.25.
Expected behavior
Report accuracy of 0.33.
Environment

torch-geometric==master

Additional context
It seems like PL has problems determining the correct batch_size of batches when data doesn't follow the conventional [batch_size, ...] format. However, it shouldn't have a problem in doing so since the batch_size can be easily inferred from the self.acc(y_hat, y_pred) call.