handle multiple inputs

pytorch · Dec 3, 2023 · 47d5916 · 47d5916
1 parent f9033f8
commit 47d5916
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchrl/objectives/value/advantages.py b/torchrl/objectives/value/advantages.py
@@ -160,7 +160,7 @@ def _call_actor_net(
     log_prob_key: NestedKey,
 ):
     # TODO: extend to handle time dimension (and vmap?)
-    log_pi = actor_net(data.select(actor_net.in_keys)).get(log_prob_key)
+    log_pi = actor_net(data.select(*actor_net.in_keys)).get(log_prob_key)
     return log_pi