[RLlib] Issue 25503: Replace torch.range with torch.arange. (#25640)

ray-project · Jun 10, 2022 · 730df43 · 730df43
1 parent b3a3519
commit 730df43
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/rllib/algorithms/dqn/dqn_torch_model.py b/rllib/algorithms/dqn/dqn_torch_model.py
@@ -148,7 +148,7 @@ def get_q_value_distributions(self, model_out):
         if self.num_atoms > 1:
             # Distributional Q-learning uses a discrete support z
             # to represent the action value distribution
-            z = torch.range(0.0, self.num_atoms - 1, dtype=torch.float32).to(
+            z = torch.arange(0.0, self.num_atoms, dtype=torch.float32).to(
                 action_scores.device
             )
             z = self.v_min + z * (self.v_max - self.v_min) / float(self.num_atoms - 1)

diff --git a/rllib/algorithms/dqn/dqn_torch_policy.py b/rllib/algorithms/dqn/dqn_torch_policy.py
@@ -62,7 +62,7 @@ def __init__(
 
         if num_atoms > 1:
             # Distributional Q-learning which corresponds to an entropy loss
-            z = torch.range(0.0, num_atoms - 1, dtype=torch.float32).to(rewards.device)
+            z = torch.arange(0.0, num_atoms, dtype=torch.float32).to(rewards.device)
             z = v_min + z * (v_max - v_min) / float(num_atoms - 1)
 
             # (batch_size, 1) * (1, num_atoms) = (batch_size, num_atoms)