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: