Loading from checkpoint giving Lookahead error

When I try to load the trainer from a checkpoint I get ‘Lookahead’ object has no attribute ‘_optimizer_load_state_dict_pre_hooks’

What do I need to do to alter the checkpoint callback to add the right pieces. I’m using lightning and I thought the checkpoint system would save everyting it needed to get going again. Here’s the cell that does the training and loading.

log_name = f’{local_config[“weights_name”]}_bands{bindices}’
log_path = r’E:/data/lsa_pytorch_logger’
logger = CSVLogger(log_path, name=log_name)

EPOCHS = 10
model.train()
torch.set_float32_matmul_precision(‘medium’)

checkpoint_callback = ModelCheckpoint(save_top_k=2, monitor=“val_F1”,mode=‘max’, dirpath = logger.log_dir, filename=‘{epoch:02d}-{val_F1:.2f}’)

trainer = pl.Trainer(accelerator=“gpu”,
devices=1,
max_epochs=EPOCHS,
logger=logger,
profiler=“simple”,
callbacks=[checkpoint_callback])

trainer.fit(model, train_loader, valid_loader)

###TO LOAD from previous CHECKPOINT
trainer.fit(model,
train_loader,
valid_loader,
ckpt_path=mpath)