Can't figure out this error with train_dataloader

[EvanZ]

Kinda losing my mind here. I have some seemingly simple code to create datasets/dataloaders from a json lines file.

class NBAMatchupDataset(Dataset):
    def __init__(self, path:str) -> None:
        dfs = []
        files = glob(path)
        for file in files:
            df = pd.read_json(file, lines=True)
            dfs.append(df)
        df = pd.concat(dfs)
        self._df = df.sample(frac=1).reset_index(drop=True)

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

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()

        row = self._df.iloc[idx]
        return {
            'pts': row['pts'],
        }

class NBAMatchupDataModule(LightningDataModule):
    def __init__(self, path:str=None, train_fraction:float=0.9, batch_size:int=32, workers:int=0) -> None:
        super().__init__()
        self._path = path
        self._train_fraction = train_fraction
        self._batch_size = batch_size
        self._workers = workers
        self.train_dataset = None
        self.val_dataset = None
        self.test_dataset = None
        
    def setup(self, stage: Optional[str]) -> None:
        dataset = NBAMatchupDataset(self._path)
        if stage=='fit':
            split_index = int(np.floor(len(dataset)*self._train_fraction))
            self.train_dataset = dataset[1:split_index-1]
            self.val_dataset = dataset[split_index:]
        if stage=='test':
            self.test_dataset = dataset
                    
    def train_dataloader(self) -> Union[DataLoader, List[DataLoader], Dict[str, DataLoader]]:
        return DataLoader(self.train_dataset, batch_size=self._batch_size, shuffle=False, pin_memory=True, \
            num_workers=self._workers)

    def val_dataloader(self) -> Union[DataLoader, List[DataLoader], Dict[str, DataLoader]]:
        return DataLoader(self.val_dataset, batch_size=self._batch_size, shuffle=False, pin_memory=True, \
            num_workers=self._workers)
    
    def test_dataloader(self) -> Union[DataLoader, List[DataLoader], Dict[str, DataLoader]]:
        return DataLoader(self.test_dataset, batch_size=self._batch_size, shuffle=False, pin_memory=True, \
            num_workers=self._workers)

Here's what I can't figure out. The test dataloader seems to work fine:

if __name__=="__main__":
    parser =ArgumentParser()
    parser.add_argument("--path", help="path/to/files", type=str, default=None)
    args = parser.parse_args()

    datamodule = NBAMatchupDataModule(args.path, train_fraction=0.9, batch_size=8)
    datamodule.setup('test')
    dataloader = datamodule.test_dataloader()
    batch = next(iter(dataloader))
    print(batch)

>> {'pts': {'home': tensor([ 5,  8,  7,  4,  1, 10,  0,  1]), 'visitor': tensor([ 2, 13,  2,  3,  3, 10,  0,  0])}}

That is the correct output. However, then I try to use the train_dataloader:

if __name__=="__main__":
    parser =ArgumentParser()
    parser.add_argument("--path", help="path/to/files", type=str, default=None)
    args = parser.parse_args()

    datamodule = NBAMatchupDataModule(args.path, train_fraction=0.9, batch_size=8)
    datamodule.setup('fit')
    dataloader = datamodule.train_dataloader()
    batch = next(iter(dataloader))
    print(batch)

This gives me the following error:

Exception has occurred: KeyError
0
  File "/home/evan/projects/***/src/datasets/nba_matchups.py", line 83, in <module>
    batch = next(iter(dataloader))

I've banged my head against this and simply can't figure out what I'm doing wrong. If anyone can spot something, I'd really appreciate it!