require 'nn'

x = torch.randn(128, 4, 84, 84)
model = nn.Sequential()

model:add(nn.SpatialConvolution(4, 16, 8, 8, 4, 4))
model:add(nn.SpatialConvolution(16, 32, 4, 4, 2, 2))
model:add(nn.SpatialConvolution(32, 256, 9, 9, 1, 1))
x_tmp = model:forward(torch.randn(1, 4, 84, 84))
n_hidden_units = x_tmp:select(1,1):numel()
print (n_hidden_units)
model:add(nn.Reshape(n_hidden_units))
model:add(nn.Linear(n_hidden_units, 4))

print (#model:forward(x))