将
def on_train_batch_start(self, batch, batch_idx, dataloader_idx):
def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
改为
def on_train_batch_start(self, batch, batch_idx):
def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):