Hey guys, need some help here. I’m trying to do an a3c agent and my implementation is based on this one[1].
My model and inputs has the same shape as [1]:
My __init__:
# num_inputs = 4
self.conv1 = nn.Conv2d(num_inputs, 32, 5, stride=1, padding=2)
self.conv2 = nn.Conv2d(32, 32, 5, stride=1, padding=1)
self.conv3 = nn.Conv2d(32, 64, 4, stride=1, padding=1)
self.conv4 = nn.Conv2d(64, 64, 3, stride=1, padding=1)

num_outputs = action_space.n # which is 4

self.lstm = nn.LSTMCell(1024, 512)
self.critic_linear = nn.Linear(512, 1)
self.actor_linear = nn.Linear(512, num_outputs)

And my forward:
inputs, (hx, cx) = inputs
# inputs is something with shape [1, 4, 80, 80]
x = F.relu(F.max_pool2d(self.conv1(inputs), kernel_size=2, stride=2))
x = F.relu(F.max_pool2d(self.conv2(x), kernel_size=2, stride=2))
x = F.relu(F.max_pool2d(self.conv3(x), kernel_size=2, stride=2))
x = F.relu(F.max_pool2d(self.conv4(x), kernel_size=2, stride=2))

x = x.view(x.size(0), -1)

hx, cx = self.lstm(x, (hx, cx))

x = hx

return self.critic_linear(x), self.actor_linear(x), (hx, cx)

I’m getting this error at  hx, cx = self.lstm(x, (hx, cx)):
RuntimeError: size mismatch, m1: [1 x 256], m2: [512 x 2048] at /py/conda-bld/pytorch_1490980628440/work/torch/lib/TH/generic/THTensorMath.c:1229
[1] https://github.com/dgriff777/rl_a3c_pytorch/blob/master/model.py 18