diff --git a/torchrl/modules/distributions/continuous.py b/torchrl/modules/distributions/continuous.py
index bae3b28c4b2..62b5df5d14b 100644
--- a/torchrl/modules/distributions/continuous.py
+++ b/torchrl/modules/distributions/continuous.py
@@ -42,9 +42,8 @@
 except ImportError:
     from torch._dynamo import is_compiling as is_dynamo_compiling
 
-TORCH_VERSION_PRE_2_6 = version.parse(torch.__version__).base_version < version.parse(
-    "2.6.0"
-)
+TORCH_VERSION = version.parse(torch.__version__).base_version
+TORCH_VERSION_PRE_2_6 = version.parse(TORCH_VERSION) < version.parse("2.6.0")
 
 
 class IndependentNormal(D.Independent):