Commit a327b27b authored by dcuai's avatar dcuai
Browse files

Update train_ddp.py

parent b3f2a8d8
...@@ -182,7 +182,7 @@ def train_model( ...@@ -182,7 +182,7 @@ def train_model(
def get_args(): def get_args():
parser = argparse.ArgumentParser(description='Train the UNet on images and target masks') parser = argparse.ArgumentParser(description='Train the UNet on images and target masks')
parser.add_argument('--local_rank', type=int, default=-1) parser.add_argument('--local-rank', type=int, default=-1)
parser.add_argument('--epochs', '-e', metavar='E', type=int, default=5, help='Number of epochs') parser.add_argument('--epochs', '-e', metavar='E', type=int, default=5, help='Number of epochs')
parser.add_argument('--batch-size', '-b', dest='batch_size', metavar='B', type=int, default=1, help='Batch size') parser.add_argument('--batch-size', '-b', dest='batch_size', metavar='B', type=int, default=1, help='Batch size')
parser.add_argument('--learning-rate', '-l', metavar='LR', type=float, default=1e-5, parser.add_argument('--learning-rate', '-l', metavar='LR', type=float, default=1e-5,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment