Running Custom validation loop causes bug when calling self.log

[vic-ene]

Hello !

I am training a model where gradients are calculated in both training and validation steps with different dataloaders.
Info: I have added the line with torch.set_grad_enabled(True) in the function validation_step to make this work.

However, I wanted to have the option of doing several passes on the validation dataloader in a row (each epoch). Thus, I am following this documentation to implement my own EvaluationLoop, but I get an error after running it with the trainer.

The class created for the EvaluationLoop is the following

from pytorch_lightning.loops import EvaluationLoop

class CustomValidationLoop(EvaluationLoop):        
    def advance(self, batch, batch_idx):
        loss = self.trainer.lightning_module.validation_step(batch, batch_idx)
        print("advanced")
        for _ in range(10):
            print("--------")

    def run(self):
        for _ in range(2):
            for i, batch in enumerate(self.dataloaders[0]):
                x, y = batch
                x = x.to(self.trainer.lightning_module.device)
                y = y.to(self.trainer.lightning_module.device)
                self.advance((x,y), i)

I assign it as such

trainer.fit_loop.epoch_loop.replace(val_loop=CustomValidationLoop)
trainer.fit(model, dl_train, dl_val)

and the error message is

File "/pytorch_lightning/core/module.py", line 390, in log
raise MisconfigurationException(
lightning_lite.utilities.exceptions.MisconfigurationException: You are trying to self.log() but it is not managed by the Trainer control flow

I use do use self.log("val_bpd", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True, sync_dist=True) in the validation_step()

The code runs normally if I do not replace the val loop.