Skip to content

Commit 6cee33e

Browse files
authored
[BugFix] PPO objective crashes if advantage_module is None (#676)
1 parent 08c0570 commit 6cee33e

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

torchrl/objectives/ppo.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,9 @@ def __init__(
8080
)
8181
self.register_buffer("gamma", torch.tensor(gamma, device=self.device))
8282
self.loss_critic_type = loss_critic_type
83-
self.advantage_module = advantage_module.to(self.device)
83+
self.advantage_module = advantage_module
84+
if self.advantage_module is not None:
85+
self.advantage_module = advantage_module.to(self.device)
8486

8587
def reset(self) -> None:
8688
pass

0 commit comments

Comments
 (0)