Fix value function?

cswinter · Dec 20, 2019 · fa0d399 · fa0d399
1 parent d829d27
commit fa0d399
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/policy_t.py b/policy_t.py
@@ -133,7 +133,7 @@ def backprop(self,
         x, x_privileged = self.latents(obs, privileged_obs)
         x = x.view(batch_size, (self.allies + self.minerals) * self.d_model)
         print(x.size())
-        values = self.value_head(x)
+        values = self.value_head(x).view(-1)
         # TODO
         #if self.use_privileged:
         #    vin = torch.cat([pooled.view(batch_size, -1), x_privileged.view(batch_size, -1)], dim=1)
@@ -188,7 +188,7 @@ def forward(self, x, x_privileged):
         #    vin = torch.cat([pooled.view(batch_size, -1), x_privileged.view(batch_size, -1)], dim=1)
         #else:
         #    vin = pooled.view(batch_size, -1)
-        values = self.value_head(x)
+        values = self.value_head(x).view(-1)
 
         logits = self.policy_head(x)
         probs = F.softmax(logits, dim=1)