diff --git a/astrai/logging.py b/astrai/logging.py
index 668ab78..996d038 100644
--- a/astrai/logging.py
+++ b/astrai/logging.py
@@ -1,11 +1,13 @@
 import logging
 import os
 
+from astrai.parallel.setup import get_rank, get_world_size
+
 
 class _DistributedContextFilter(logging.Filter):
     def filter(self, record: logging.LogRecord) -> bool:
-        record.rank = os.environ.get("RANK", "0")
-        record.world_size = os.environ.get("WORLD_SIZE", "1")
+        record.rank = str(get_rank())
+        record.world_size = str(get_world_size())
         return True
 
 
diff --git a/astrai/parallel/setup.py b/astrai/parallel/setup.py
index e79eacb..4709d75 100644
--- a/astrai/parallel/setup.py
+++ b/astrai/parallel/setup.py
@@ -30,15 +30,13 @@ def get_current_device():
 def get_world_size() -> int:
     if dist.is_available() and dist.is_initialized():
         return dist.get_world_size()
-    else:
-        return 1
+    return int(os.environ.get("WORLD_SIZE", "1"))
 
 
 def get_rank() -> int:
     if dist.is_available() and dist.is_initialized():
         return dist.get_rank()
-    else:
-        return 0
+    return int(os.environ.get("RANK", "0"))
 
 
 @contextmanager