add max steps
parent
055e7ee48a
commit
0a19a59564
|
@ -337,7 +337,8 @@ if __name__ == '__main__':
|
|||
parser.add_argument('--work-dir', '-o', type=str, default='out', help="output working directory")
|
||||
parser.add_argument('--resume', action='store_true', help="when this flag is used, we will resume optimization from existing model in the workdir")
|
||||
parser.add_argument('--sample-only', action='store_true', help="just sample from the model and quit, don't train")
|
||||
parser.add_argument('--num-workers', '-n', type=int, default=0, help="number of data workers for both train/test")
|
||||
parser.add_argument('--num-workers', '-n', type=int, default=4, help="number of data workers for both train/test")
|
||||
parser.add_argument('--max-steps', type=int, default=-1, help="max number of optimization steps to run for, or -1 for infinite.")
|
||||
parser.add_argument('--device', type=str, default='cpu', help="device to use for compute, examples: cpu|cuda|cuda:2|mps")
|
||||
parser.add_argument('--seed', type=int, default=3407, help="seed")
|
||||
# sampling
|
||||
|
@ -439,3 +440,7 @@ if __name__ == '__main__':
|
|||
print_samples(num=10)
|
||||
|
||||
step += 1
|
||||
# termination conditions
|
||||
if args.max_steps >= 0 and step >= args.max_steps:
|
||||
break
|
||||
|
||||
|
|
Loading…
Reference in New Issue