From c12b01701bf1bf62fe5b124d5a410c2fdc524a8c Mon Sep 17 00:00:00 2001 From: Xuehai Pan Date: Sun, 13 Aug 2023 22:15:56 +0800 Subject: [PATCH] fix(callbacks): fix a potential attribute errors for newer PyTorch-Lightning --- docs/source/spelling_wordlist.txt | 1 + nvitop/callbacks/pytorch_lightning.py | 6 +++++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/docs/source/spelling_wordlist.txt b/docs/source/spelling_wordlist.txt index d75fce7..106df4f 100644 --- a/docs/source/spelling_wordlist.txt +++ b/docs/source/spelling_wordlist.txt @@ -150,3 +150,4 @@ KiB tx rx ThroughputInfo +pytorch diff --git a/nvitop/callbacks/pytorch_lightning.py b/nvitop/callbacks/pytorch_lightning.py index 0c49ce4..8af3f56 100644 --- a/nvitop/callbacks/pytorch_lightning.py +++ b/nvitop/callbacks/pytorch_lightning.py @@ -119,7 +119,11 @@ class GpuStatsLogger(Callback): # pylint: disable=too-many-instance-attributes f'The root device type is {trainer.strategy.root_device.type}.', ) - device_ids = trainer.data_parallel_device_ids + try: + device_ids = trainer.device_ids # pytorch-lightning >= 1.6.0 + except AttributeError: + device_ids = trainer.data_parallel_device_ids # pytorch-lightning < 1.6.0 + try: self._devices = get_devices_by_logical_ids(device_ids, unique=True) except (libnvml.NVMLError, RuntimeError) as ex: