diff --git a/domainlab/arg_parser.py b/domainlab/arg_parser.py
index 6812c0356..1e792fa88 100644
--- a/domainlab/arg_parser.py
+++ b/domainlab/arg_parser.py
@@ -53,6 +53,10 @@ def mk_parser_main():
 
     parser.add_argument('--epo_te', default=1, type=int,
                         help='test performance per {} epochs')
+    
+    parser.add_argument('-w', '--warmup', type=int, default=100,
+                        help='number of epochs for hyper-parameter warm-up. \
+                        Set to 0 to turn warmup off.')
 
     parser.add_argument('--debug', action='store_true', default=False)
     parser.add_argument('--dmem', action='store_true', default=False)
diff --git a/domainlab/models/args_vae.py b/domainlab/models/args_vae.py
index 5fd22aa69..7ac8ae9a0 100644
--- a/domainlab/models/args_vae.py
+++ b/domainlab/models/args_vae.py
@@ -46,9 +46,4 @@ def add_args2parser_vae(parser):
                         help='multiplier for KL x')
     parser.add_argument('--beta_y', type=float, default=1.,
                         help='multiplier for KL y')
-    #
-
-    parser.add_argument('-w', '--warmup', type=int, default=100,
-                        help='number of epochs for hyper-parameter warm-up. \
-                        Set to 0 to turn warmup off.')
     return parser