Update train.py

splendor1811 · May 18, 2023 · f674ad9 · f674ad9
1 parent 23ff2c8
commit f674ad9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/train.py b/train.py
@@ -6,7 +6,7 @@
 from callbacks.callbacks import get_callbacks, cosine_annealing_with_warmup
 from dataloader.dataloader import build_augmenter, build_dataset, build_decoder
 # from supervision.dataloader import build_augmenter, build_dataset, build_decoder
-from model import create_segment_model, build_model
+from model import build_model
 import os 
 import tensorflow_addons as tfa 
 from optimizers.lion_opt import Lion
@@ -123,4 +123,4 @@ def myprint(s):
 model.load_weights(save_path)
 
 model.evaluate(test_dataset)
-model.save("final_model.h5")
+model.save("final_model.h5")