Merge c8000de into b9d5f5e

Lightning-Universe · Jan 30, 2021 · cb1aa65 · cb1aa65
2 parents b9d5f5e + c8000de
commit cb1aa65
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/pl_bolts/models/rl/double_dqn_model.py b/pl_bolts/models/rl/double_dqn_model.py
@@ -56,7 +56,7 @@ def training_step(self, batch: Tuple[torch.Tensor, torch.Tensor], _) -> OrderedD
  """
 
  # calculates training loss
- loss = double_dqn_loss(batch, self.net, self.target_net)
+ loss = double_dqn_loss(batch, self.net, self.target_net, self.gamma)
 
  if self.trainer.use_dp or self.trainer.use_ddp2:
  loss = loss.unsqueeze(0)

diff --git a/pl_bolts/models/rl/dqn_model.py b/pl_bolts/models/rl/dqn_model.py
@@ -270,7 +270,7 @@ def training_step(self, batch: Tuple[torch.Tensor, torch.Tensor], _) -> OrderedD
  """
 
  # calculates training loss
- loss = dqn_loss(batch, self.net, self.target_net)
+ loss = dqn_loss(batch, self.net, self.target_net, self.gamma)
 
  if self.trainer.use_dp or self.trainer.use_ddp2:
  loss = loss.unsqueeze(0)

diff --git a/pl_bolts/models/rl/per_dqn_model.py b/pl_bolts/models/rl/per_dqn_model.py
@@ -114,7 +114,7 @@ def training_step(self, batch, _) -> OrderedDict:
  indices = indices.cpu().numpy()
 
  # calculates training loss
- loss, batch_weights = per_dqn_loss(samples, weights, self.net, self.target_net)
+ loss, batch_weights = per_dqn_loss(samples, weights, self.net, self.target_net, self.gamma)
 
  if self.trainer.use_dp or self.trainer.use_ddp2:
  loss = loss.unsqueeze(0)