From 1c7635556f131f502d5af342012974d63ea483c6 Mon Sep 17 00:00:00 2001
From: Colin <>
Date: Sat, 9 Aug 2025 16:32:48 +0800
Subject: [PATCH] Reconfig train model and code.

---
 wit/dataset/meaning_dataset.py | 2 +-
 wit/train.py                   | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/wit/dataset/meaning_dataset.py b/wit/dataset/meaning_dataset.py
index f637b62..df0cfcc 100644
--- a/wit/dataset/meaning_dataset.py
+++ b/wit/dataset/meaning_dataset.py
@@ -6,7 +6,7 @@ from typing import Dict, Tuple
 from torch.utils.data import ConcatDataset, DataLoader, Dataset, random_split
 import numpy as np
 from torch.utils.data import BatchSampler
-from node_tree import NodeTree
+from dataset.node_tree import NodeTree
 
 
 class MeaningMap:
diff --git a/wit/train.py b/wit/train.py
index 9021e32..48a6c99 100644
--- a/wit/train.py
+++ b/wit/train.py
@@ -39,7 +39,7 @@ if __name__ == "__main__":
     config.vocab_size = 32
     config.hidden_size = 256  # 128 1024 2048  32
     config.intermediate_size = 512
-    config.num_hidden_layers = 4  # 6 12 24  3
+    config.num_hidden_layers = 3  # 6 12 24  3
     config.num_attention_heads = 4  # 8 8 16
 
     torch.manual_seed(conf.seed)