From 062a27b83fd7054956bcbc8155b9f87289d58337 Mon Sep 17 00:00:00 2001
From: mhz <cxyoz@outlook.com>
Date: Wed, 26 Jun 2024 22:10:07 +0200
Subject: [PATCH] try update the api in DataInfo

---
 graph_dit/main.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/graph_dit/main.py b/graph_dit/main.py
index d7fbe14..24e3c39 100644
--- a/graph_dit/main.py
+++ b/graph_dit/main.py
@@ -78,7 +78,7 @@ def main(cfg: DictConfig):
 
     datamodule = dataset.DataModule(cfg)
     datamodule.prepare_data()
-    dataset_infos = dataset.DataInfos(datamodule=datamodule, cfg=cfg)
+    dataset_infos = dataset.DataInfos(datamodule=datamodule, cfg=cfg, dataset=datamodule.dataset)
     # train_smiles, reference_smiles = datamodule.get_train_smiles()
 
     # get input output dimensions