Update train.py

train.py

@@ -10,7 +10,7 @@ from pytorch_lightning.callbacks import ModelCheckpoint
 from torch.utils.data import DataLoader
 from huggingface_hub import PyTorchModelHubMixin
 
-from isnet import ISNetDIS
+from isnet import ISNetDIS, ISNetGTEncoder
 
 
 # warnings.filterwarnings("ignore")
@@ -22,6 +22,8 @@ def get_net(net_name, img_size):
         return ISNetDIS()
     elif net_name == "isnet_is":
         return ISNetDIS()
+    elif net_name == "isnet_gt":
+        return ISNetGTEncoder()
     raise NotImplementedError