diff --git a/diffusers_trainer.py b/diffusers_trainer.py
index b51c317..10eef02 100644
--- a/diffusers_trainer.py
+++ b/diffusers_trainer.py
@@ -81,6 +81,7 @@ parser.add_argument('--image_log_inference_steps', type=int, default=50, help='N
 parser.add_argument('--image_log_scheduler', type=str, default="PNDMScheduler", help='Number of inference steps to use to log images.')
 parser.add_argument('--clip_penultimate', type=bool, default=False, help='Use penultimate CLIP layer for text embedding')
 parser.add_argument('--output_bucket_info', type=bool, default=False, help='Outputs bucket information and exits')
+parser.add_argument('--use_xformers', action='store_true',help='Use memory efficient attention')
 args = parser.parse_args()
 
 def setup():
@@ -549,6 +550,9 @@ def main():
 
     if args.gradient_checkpointing:
         unet.enable_gradient_checkpointing()
+    
+    if args.use_xformers:
+        unet.set_use_memory_efficient_attention_xformers(True)
 
     if args.use_8bit_adam: # Bits and bytes is only supported on certain CUDA setups, so default to regular adam if it fails.
         try: