Update train.py

train.py

@@ -10,8 +10,7 @@ from pytorch_lightning.callbacks import ModelCheckpoint
 from torch.utils.data import DataLoader
 from huggingface_hub import PyTorchModelHubMixin
 
-from model import ISNetDIS, ISNetGTEncoder, U2NET, U2NET_full2, U2NET_lite2, MODNet \
-    , InSPyReNet, InSPyReNet_Res2Net50, InSPyReNet_SwinB
+from isnet import ISNetDIS
 
 
 # warnings.filterwarnings("ignore")
@@ -23,18 +22,6 @@ def get_net(net_name, img_size):
         return ISNetDIS()
     elif net_name == "isnet_is":
         return ISNetDIS()
-    elif net_name == "isnet_gt":
-        return ISNetGTEncoder()
-    elif net_name == "u2net":
-        return U2NET_full2()
-    elif net_name == "u2netl":
-        return U2NET_lite2()
-    elif net_name == "modnet":
-        return MODNet()
-    elif net_name == "inspyrnet_res":
-        return InSPyReNet_Res2Net50(base_size=img_size)
-    elif net_name == "inspyrnet_swin":
-        return InSPyReNet_SwinB(base_size=img_size)
     raise NotImplementedError