Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Update test_pruning.py to use `devices` instead of `gpus` or `ipus` (#11339)

This commit is contained in:
Jv Kyle Eclarin 2022-01-08 10:24:29 -06:00 committed by GitHub
parent 4710a8128b
commit b56d8677ad
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
tests/callbacks/test_device_stats_monitor.py
View File

@ -44,7 +44,8 @@ def test_device_stats_gpu_from_torch(tmpdir):
max_epochs=2,
limit_train_batches=7,
log_every_n_steps=1,
gpus=1,
accelerator="gpu",
devices=1,
callbacks=[device_stats],
logger=DebugLogger(tmpdir),
enable_checkpointing=False,
@ -73,7 +74,8 @@ def test_device_stats_gpu_from_nvidia(tmpdir):
max_epochs=2,
limit_train_batches=7,
log_every_n_steps=1,
gpus=1,
accelerator="gpu",
devices=1,
callbacks=[device_stats],
logger=DebugLogger(tmpdir),
enable_checkpointing=False,
@ -101,7 +103,8 @@ def test_device_stats_monitor_tpu(tmpdir):
default_root_dir=tmpdir,
max_epochs=1,
limit_train_batches=1,
tpu_cores=8,
accelerator="tpu",
devices=8,
log_every_n_steps=1,
callbacks=[device_stats],
logger=DebugLogger(tmpdir),