Update

[ghstack-poisoned]
pytorch · Sep 13, 2024 · 00b1f8f · 00b1f8f
2 parents 7b79cbb + 0bc468c
commit 00b1f8f
Show file tree

Hide file tree

Showing 8 changed files with 17 additions and 10 deletions.
diff --git a/.github/unittest/linux/scripts/run_all.sh b/.github/unittest/linux/scripts/run_all.sh
@@ -127,13 +127,13 @@ if [[ "$TORCH_VERSION" == "nightly" ]]; then
   if [ "${CU_VERSION:-}" == cpu ] ; then
       pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu -U
   else
-      pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/$CU_VERSION
+      pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/$CU_VERSION -U
   fi
 elif [[ "$TORCH_VERSION" == "stable" ]]; then
     if [ "${CU_VERSION:-}" == cpu ] ; then
-      pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+      pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cpu -U
   else
-      pip3 install torch torchvision --index-url https://download.pytorch.org/whl/$CU_VERSION
+      pip3 install torch torchvision --index-url https://download.pytorch.org/whl/$CU_VERSION -U
   fi
 else
   printf "Failed to install pytorch"

diff --git a/.github/unittest/linux_libs/scripts_brax/install.sh b/.github/unittest/linux_libs/scripts_brax/install.sh
@@ -34,7 +34,7 @@ if [[ "$TORCH_VERSION" == "nightly" ]]; then
   fi
 elif [[ "$TORCH_VERSION" == "stable" ]]; then
     if [ "${CU_VERSION:-}" == cpu ] ; then
-      pip3 install torch --index-url https://download.pytorch.org/whl/cpu
+      pip3 install torch --index-url https://download.pytorch.org/whl/cpu -U
   else
       pip3 install torch --index-url https://download.pytorch.org/whl/cu121
   fi

diff --git a/.github/unittest/linux_libs/scripts_openx/install.sh b/.github/unittest/linux_libs/scripts_openx/install.sh
@@ -37,9 +37,9 @@ if [[ "$TORCH_VERSION" == "nightly" ]]; then
   fi
 elif [[ "$TORCH_VERSION" == "stable" ]]; then
     if [ "${CU_VERSION:-}" == cpu ] ; then
-      pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+      pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cpu -U
   else
-      pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cu121
+      pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cu121 -U
   fi
 else
   printf "Failed to install pytorch"

diff --git a/.github/unittest/linux_libs/scripts_vd4rl/install.sh b/.github/unittest/linux_libs/scripts_vd4rl/install.sh
@@ -37,7 +37,7 @@ if [[ "$TORCH_VERSION" == "nightly" ]]; then
   fi
 elif [[ "$TORCH_VERSION" == "stable" ]]; then
     if [ "${CU_VERSION:-}" == cpu ] ; then
-      pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+      pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cpu -U
   else
       pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cu121
   fi

diff --git a/.github/unittest/linux_optdeps/scripts/install.sh b/.github/unittest/linux_optdeps/scripts/install.sh
@@ -20,7 +20,7 @@ version="$(python -c "print('.'.join(\"${CUDA_VERSION}\".split('.')[:2]))")"
 git submodule sync && git submodule update --init --recursive
 
 printf "Installing PyTorch with %s\n" "${CU_VERSION}"
-pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/$CU_VERSION
+pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/$CU_VERSION -U
 
 # install tensordict
 if [[ "$RELEASE" == 0 ]]; then

diff --git a/test/test_cost.py b/test/test_cost.py
@@ -146,6 +146,7 @@
     _split_and_pad_sequence,
 )
 
+TORCH_VERSION = torch.__version__
 
 # Capture all warnings
 pytestmark = [
@@ -15644,6 +15645,7 @@ def __init__(self):
                 assert p.device == dest
 
 
+@pytest.mark.skipif(TORCH_VERSION < "2.5", reason="requires torch>=2.5")
 def test_exploration_compile():
     m = ProbabilisticTensorDictModule(
         in_keys=["loc", "scale"],

diff --git a/torchrl/modules/distributions/continuous.py b/torchrl/modules/distributions/continuous.py
@@ -12,7 +12,12 @@
 import numpy as np
 import torch
 from torch import distributions as D, nn
-from torch.compiler import assume_constant_result
+
+try:
+    from torch.compiler import assume_constant_result
+except ImportError:
+    from torch._dynamo import assume_constant_result
+
 from torch.distributions import constraints
 from torch.distributions.transforms import _InverseTransform
 

diff --git a/torchrl/objectives/common.py b/torchrl/objectives/common.py
@@ -26,7 +26,7 @@
 
 try:
     from torch.compiler import is_dynamo_compiling
-except ImportError:
+except ModuleNotFoundError:
     from torch._dynamo import is_compiling as is_dynamo_compiling