I have the following dqn model:

```
class DQN(nn.Module):
def __init__(self, input_shape, n_actions):
super(DQN, self).__init__()
self.conv = nn.Sequential(
nn.Conv2d(input_shape[0], 32, kernel_size=8, stride=4),
nn.ReLU(),
nn.Conv2d(32, 64, kernel_size=4, stride=2),
nn.ReLU(),
nn.Conv2d(64, 64, kernel_size=3, stride=1),
nn.ReLU()
)
conv_out_size = self._get_conv_out(input_shape)
self.fc = nn.Sequential(
nn.Linear(conv_out_size, 512),
nn.ReLU(),
nn.Linear(512, n_actions)
)
def _get_conv_out(self, shape):
o = self.conv(Variable(torch.zeros(1, *shape)))
return int(np.prod(o.size()))
def forward(self, x):
fx = x.float() / 256.0
conv_out = self.conv(fx).view(fx.size()[0], -1)
return self.fc(conv_out)
```

I trained it for Pong and got more than 18.0 mean reward for the last 100 games.

It was saved:

```
torch.save(net.state_dict(), 'pong_model.pt')
```

But when I try to load it and play some games I get wrong results.

```
net = models.DQN(env.observation_space.shape, env.action_space.n)
net.load_state_dict(torch.load('pong_model.pt', map_location=lambda storage, loc: storage))
for i in range(20):
state = env.reset()
while True:
env.render()
action = agent(state, net)
next_state, _, done, _ = env.step(action)
if done:
break
state = next_state
```

It can’t win even once. What could be wrong?