[BugFix] Torch 1.13 compat (#1294)

pytorch · Jun 17, 2023 · 3b8e128 · 3b8e128 · github-actions · Jun 17, 2023
1 parent 593961e
commit 3b8e128
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 4 deletions.
diff --git a/torchrl/objectives/value/advantages.py b/torchrl/objectives/value/advantages.py
@@ -34,6 +34,16 @@
     vec_td_lambda_return_estimate,
 )
 
+try:
+    from torch import vmap
+except ImportError as err:
+    try:
+        from functorch import vmap
+    except ImportError:
+        raise ImportError(
+            "vmap couldn't be found. Make sure you have torch>1.13 installed."
+        ) from err
+
 
 def _self_set_grad_enabled(fun):
     @wraps(fun)
@@ -111,9 +121,9 @@ def _call_value_nets(
             )
         elif params is not None:
             params_stack = torch.stack([params, next_params], 0)
-            data_out = torch.vmap(value_net, (0, 0))(data_in, params_stack)
+            data_out = vmap(value_net, (0, 0))(data_in, params_stack)
         else:
-            data_out = torch.vmap(value_net, (0,))(data_in)
+            data_out = vmap(value_net, (0,))(data_in)
         value_est = data_out.get(value_key)
         value, value_ = value_est[0], value_est[1]
     data.set(value_key, value)

diff --git a/tutorials/sphinx-tutorials/torchrl_demo.py b/tutorials/sphinx-tutorials/torchrl_demo.py
@@ -474,10 +474,10 @@
 
 ###############################################################################
 
-import functorch
+from torch import vmap
 
 params_expand = params.expand(4)
-tensordict_exp = functorch.vmap(sequence, (None, 0))(tensordict, params_expand)
+tensordict_exp = vmap(sequence, (None, 0))(tensordict, params_expand)
 print(tensordict_exp)
 
 ###############################################################################