Merge pull request #159 from WilfChen/c51-range-error-fix

fix c51 range type error
mindspore-lab · Jan 13, 2024 · 7fea931 · 7fea931
2 parents 4734caa + 848f03f
commit 7fea931
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/mindspore_rl/algorithm/c51/c51.py b/mindspore_rl/algorithm/c51/c51.py
@@ -241,10 +241,9 @@ def next_distribution(self, next_observation, batch_size):
         next_action = self.get_max_index(next_target_q_values)[0]
         next_qt_argmax = self.expand_dims(next_action, 1)
         next_qt_argmax = self.cast(next_qt_argmax, ms.int32)
-        batch_indices = self.get_range(
-            Tensor(0, ms.int32), Tensor(batch_size, ms.int32), Tensor(1, ms.int32)
-        )
+        batch_indices = self.get_range(0, batch_size, 1)
         batch_indices = self.expand_dims(batch_indices, 1).reshape(batch_size, 1)
+        batch_indices = ops.cast(batch_indices, ms.int32)
         next_qt_index = self.concat((batch_indices, next_qt_argmax))
         return self.gather_nd(next_target_probabilities, next_qt_index)