Update tests/plugins/test_sharded_plugin.py

Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>
This commit is contained in:
Sean Naren 2020-11-26 16:37:22 +00:00 committed by GitHub
parent 091c236392
commit ff34a8fed9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 2 deletions

View File

@ -458,8 +458,9 @@ def run_sharded_correctness(
max_percent_speed_diff=0.25,
model_cls=TestModel):
"""
Ensures that the trained model is identical to the standard DDP implementation.
Also checks for speed/memory regressions, we should expect always less memory but performance to fluctuate.
Ensures that the trained model is identical to the standard DDP implementation.
Also checks for speed/memory regressions, we should expect always less memory but performance to fluctuate.
Args:
accelerator: Accelerator type for test.
gpus: Number of GPUS to enable.