diff --git a/model_helpers/prepare_positionnet.py b/model_helpers/prepare_positionnet.py
index fdd55f5..4eff4e6 100644
--- a/model_helpers/prepare_positionnet.py
+++ b/model_helpers/prepare_positionnet.py
@@ -9,13 +9,13 @@ def get_positionnet_default_params():
         "mid_dim": 3072,
         "out_dim": 768,
         "test_drop_boxes": False,
-        "test_drop_masks": False,
+        "test_drop_masks": True,
         "test_drop_points": False,
-        "test_drop_scribbles": False,
+        "test_drop_scribbles": True,
         "train_add_boxes": True,
-        "train_add_masks": False,
-        "train_add_points": False,
-        "train_add_scribbles": False,
+        "train_add_masks": True,
+        "train_add_points": True,
+        "train_add_scribbles": True,
         "use_seperate_tokenizer": True,
     }
 
diff --git a/modules/text_grounding_net.py b/modules/text_grounding_net.py
index ba77d7f..b84d057 100644
--- a/modules/text_grounding_net.py
+++ b/modules/text_grounding_net.py
@@ -26,7 +26,8 @@ class UniFusion(nn.Module):
         self.use_seperate_tokenizer = use_seperate_tokenizer
 
         # Use instance masks as additional model inputs for mask conditioned image generation
-        self.use_segs = True if self.add_masks else False
+        #self.use_segs = True if self.add_masks else False
+        self.use_segs = False
 
         if self.use_segs:
             in_dim = 30