diff --git a/models/pretraining.py b/models/pretraining.py
index fa25741..e824e17 100644
--- a/models/pretraining.py
+++ b/models/pretraining.py
@@ -52,6 +52,7 @@ def __init__(self, args):
             num_layers=args.graph_num_layers,
             num_heads=args.graph_num_heads,
             dropout=args.dropout,
+            attention_dropout=getattr(args, "attention_dropout", args.dropout),
             dim_feedforward=args.dim_feedforward,
             add_positional_encoding=args.add_positional_encoding,
             use_edge_bias=args.use_edge_bias,
diff --git a/pretrain.py b/pretrain.py
index 372459d..e5d4767 100644
--- a/pretrain.py
+++ b/pretrain.py
@@ -26,6 +26,7 @@ def build_parser():
     parser.add_argument("--dim_feedforward", type=int, default=512)
     parser.add_argument("--dropout", type=float, default=0.25)
     parser.add_argument("--mlp_dropout", type=float, default=0.1)
+    parser.add_argument("--attention_dropout", type=float, default=0.1)
     parser.add_argument("--act", type=str, default="gelu")
     parser.add_argument("--curve_emb_dim", type=int, default=64)
     parser.add_argument("--surface_emb_dim", type=int, default=64)