diff --git a/data/data.py b/data/data.py
index 672daa5..e0b8a53 100644
--- a/data/data.py
+++ b/data/data.py
@@ -88,6 +88,13 @@ def test_dataset():
         logger.info(f"Split: {split}")
         
         # ... (其余测试代码保持不变) ...
+        dataset = BRepSDFDataset(brep_data_dir='/home/wch/myDeepSDF/test_data/pkl', sdf_data_dir='/home/wch/myDeepSDF/test_data/sdf', split='train')
+        dataloader = torch.utils.data.DataLoader(dataset, batch_size=4, shuffle=True)
+
+        for batch in dataloader:
+            print(batch['brep_features'].shape)
+            print(batch['sdf'].shape)
+            break
 
 if __name__ == '__main__':
     test_dataset()
\ No newline at end of file