From eaee962fa9e50a814a288208f1eee749f3c4e8b6 Mon Sep 17 00:00:00 2001
From: LinB203 <2267330597@qq.com>
Date: Fri, 8 Mar 2024 17:25:48 +0800
Subject: [PATCH] update train.py

Former-commit-id: 844173711342f2dcc021294e87971ca39854bfe4
---
 opensora/train/train.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/opensora/train/train.py b/opensora/train/train.py
index 7c7c0a9f5..7c908e85c 100644
--- a/opensora/train/train.py
+++ b/opensora/train/train.py
@@ -50,6 +50,8 @@ def main(args):
     # Setup accelerator:
     accelerator = Accelerator()
     device = accelerator.device
+    args.mixed_precision = accelerator.mixed_precision
+    args.gradient_accumulation_steps = accelerator.gradient_accumulation_steps
 
     # Setup an experiment folder:
     if accelerator.is_main_process:
@@ -314,11 +316,9 @@ if __name__ == "__main__":
     parser.add_argument("--dynamic-frames", action="store_true")
     parser.add_argument("--resume-from-checkpoint", action="store_true")
     parser.add_argument("--gradient-checkpointing", action="store_true")
-    parser.add_argument("--mixed-precision", action="store_true")
     parser.add_argument("--use-compile", action="store_true")
     parser.add_argument("--lr", type=float, default=1e-4)
     parser.add_argument("--lr-warmup-steps", type=int, default=0)
-    parser.add_argument("--gradient-accumulation-steps", type=int, default=0)
 
     parser.add_argument("--clip-grad-norm", default=None, type=float, help="the maximum gradient norm (default None)")
     # --------------------------------------