diff --git a/models/roma_unsb_model.py b/models/roma_unsb_model.py
index 046cbdd..b4df892 100644
--- a/models/roma_unsb_model.py
+++ b/models/roma_unsb_model.py
@@ -218,6 +218,8 @@ class RomaUnsbModel(BaseModel):
         
         parser.add_argument('--lambda_inc', type=float, default=1.0, help='incremental weight for content-aware optimization')
         parser.add_argument('--eta_ratio', type=float, default=0.1, help='ratio of content-rich regions')
+
+        parser.add_argument('--atten_layers', type=str, default='1,3,5', help='compute Cross-Similarity on which layers')
         
 
         parser.set_defaults(pool_size=0)  # no image pooling