Update main.py

ZhuLvs · Aug 28, 2024 · 25639a4 · 25639a4
1 parent 2dc052d
commit 25639a4
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/model/main.py b/model/main.py
@@ -3,6 +3,7 @@
 import torch.optim as optim
 from torch.utils.tensorboard import SummaryWriter
 from model import CustomDeepLabV3
+#from feature_fusion_model import CustomDeepLabV3
 from Dataset import CustomDataset, transform, val_transform
 from Maskloss import MaskedL1Loss
 from Inference import train_one_epoch, evaluate
@@ -45,4 +46,4 @@
 
 
     writer.close()
-    torch.save(model.state_dict(), final_model_save_path)
+    torch.save(model.state_dict(), final_model_save_path)