From e7e94d715bece1f349bc8b3d85b085100f85dfb9 Mon Sep 17 00:00:00 2001 From: Shangtong Zhang Date: Sat, 31 Mar 2018 21:46:16 -0600 Subject: [PATCH] Add keepdim --- network/base_network.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/network/base_network.py b/network/base_network.py index 3c3e143..075f86e 100644 --- a/network/base_network.py +++ b/network/base_network.py @@ -88,7 +88,7 @@ class DuelingNet(BasicNet): phi = self.forward(x) value = self.fc_value(phi) advantange = self.fc_advantage(phi) - q = value.expand_as(advantange) + (advantange - advantange.mean(1).expand_as(advantange)) + q = value.expand_as(advantange) + (advantange - advantange.mean(1, keepdim=True).expand_as(advantange)) if to_numpy: return q.cpu().data.numpy() return q