What causes the Pytorch nn Module to return a 1 or 0 for the argmax of this array?
I have an array containing four random numbers that an argmax function should be returning a 0, 1, 2, or 3, but when the argmax is called from inside the nn.Module model it is always 0 or 1.
I would just like to know how and why it is always getting a 1 or 0 from the four numbers in the array.
Below I have the nn.Module and a comparison of a random array of len 3 computed inside and outside of the model using the function act (Net.act).
from torch import nn
import torch

import random


class Network(nn.Module):
    def __init__(self):
        super().__init__()

        self.net = nn.Sequential(
            nn.Linear(4, 64),
            nn.Tanh(),
            nn.Linear(64, 2))

    def forward(self, x):
        return self.net(x)

    def act(self, obs):
        obs_t = torch.as_tensor(obs, dtype=torch.float32, device=device)
        q_values = self(obs_t.unsqueeze(0))
        max_q_index = torch.argmax(q_values, dim=1)[0]
        action = max_q_index.detach().item()

        return action


Net = Network()

This is a side by side comparison of a random array of len 3 inside and outside of the act function in the nn.Module (Net).
for _ in range(20):
    z = np.array([random.uniform(-1, 1) for _ in range(4)])

    obs_t = torch.as_tensor(z, dtype=torch.float32, device=device)
    q_values = (obs_t.unsqueeze(0))
    max_q_index = torch.argmax(q_values, dim=1)[0]
    action = max_q_index.detach().item()

    print(action, Net.act(z))

Output
3 0
0 0
2 0
0 0
2 0
0 0
1 0
2 0
3 0
0 1
0 0
1 1
3 1
3 0
1 0
0 1
3 0
3 0
0 0
1 0


self(obs_t.unsqueeze(0)) returns a 2 column matrix because the last layer of your model (nn.Linear(64, 2)) is defined to output two columns. max_q_index contains the column index of the largest value in each row of the model output (column index because dim=1). Since there are only 2 columns, max_q_index can only have values of 0 or 1.
