fixed a minor bug

816fb890 · Mostofa Patwary · b2f57fc4 · 816fb890
Commit 816fb890 authored Feb 17, 2021 by Mostofa Patwary
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

megatron/training.py megatron/training.py +1 -1

No files found.
--- a/megatron/training.py
+++ b/megatron/training.py
@@ -311,7 +311,7 @@ def setup_model_and_optimizer(model_provider_func):
    # We only support local DDP with multiple micro-batches.
    if get_num_microbatches() > 1:
        assert args.DDP_impl == 'local'
-    if len(model) == 1:
+    if len(model) > 1:
        assert args.DDP_impl == 'local'
    if mpu.get_pipeline_model_parallel_world_size() > 1:
        assert args.DDP_impl == 'local'