From d6c7c7c6755029d03dbc6f286bd0a749ecd828ae Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Wed, 4 Nov 2020 18:04:03 +0100 Subject: [PATCH] [RLlib] Make sure, DQN torch actions are of type=long before torch.nn.functional.one_hot() op. (#11800) --- rllib/agents/dqn/dqn_torch_policy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rllib/agents/dqn/dqn_torch_policy.py b/rllib/agents/dqn/dqn_torch_policy.py index dbaa9a938..0d633fa06 100644 --- a/rllib/agents/dqn/dqn_torch_policy.py +++ b/rllib/agents/dqn/dqn_torch_policy.py @@ -253,7 +253,7 @@ def build_q_losses(policy: Policy, model, _, is_training=True) # Q scores for actions which we know were selected in the given state. - one_hot_selection = F.one_hot(train_batch[SampleBatch.ACTIONS], + one_hot_selection = F.one_hot(train_batch[SampleBatch.ACTIONS].long(), policy.action_space.n) q_t_selected = torch.sum( torch.where(q_t > FLOAT_MIN, q_t,