def METHOD_NAME(
    self,
    trainer,
    pl_module: AtomisticTask,
    prediction: Any,
    batch_indices: List[int],
    batch: Any,
    batch_idx: int,
    dataloader_idx: int,
):
    bdir = os.path.join(self.output_dir, str(dataloader_idx))
    os.makedirs(bdir, exist_ok=True)
    torch.save(prediction, os.path.join(bdir, f"{batch_idx}.pt"))