Index 5 is out of bounds for dimension 0 with size 3

Bocksa: <p>Hi, I’ve been encountering this issue for the past day and have been unsure as to how to fix it. The index which is out of bounds isn’t consistent either.</p>
<pre data-code-wrap="py"><code class="lang-py">import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import os

class Linear_QNet(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        self.linear1 = nn.Linear(input_size, hidden_size)
        self.linear2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = F.relu(self.linear1(x))
        x = self.linear2(x)
        return x
    
    def save(self, file_name='model.pth'):
        model_folder_path = './model'
        if not os.path.exists(model_folder_path):
            os.makedirs(model_folder_path)

        file_name = os.path.join(model_folder_path, file_name)
        torch.save(self.state_dict(), file_name)

class QTrainer:
    def __init__(self, model, lr, gamma):
        self.lr = lr
        self.gamma = gamma
        self.model = model
        self.optimiser = optim.Adam(model.parameters(), lr=self.lr)
        self.criterion = nn.MSELoss() # Loss function

    def train_step(self, state, action, reward, next_state, done):
        state = torch.tensor(state, dtype=torch.float)
        next_state = torch.tensor(next_state, dtype=torch.float)
        action = torch.tensor(action, dtype=torch.long)
        reward = torch.tensor(reward, dtype=torch.float)

        if len(state.shape) == 1:
            state = torch.unsqueeze(state, 0)
            next_state = torch.unsqueeze(next_state, 0)
            action = torch.unsqueeze(action, 0)
            reward = torch.unsqueeze(reward, 0)
            done = (done, )

        pred = self.model(state)
        target = pred.clone()

        for i in range(len(done) - 1):
            Q_new = reward[i]
            
            if not done[i]:
                Q_new = reward[i] + self.gamma * torch.max(self.model(next_state[i]))
            
            targmax = torch.argmax(action).item()
            target[i][targmax] = Q_new
        
        self.optimiser.zero_grad()
        loss = self.criterion(target, pred)
        loss.backward()
        self.optimiser.step()
</code></pre>
<p><code>target[i][targmax] = Q_new</code> is the line which has the issue</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/d/6d0d5f2696fa6a513e075fcb2bf76d89af915f96.png" data-download-href="https://discuss.pytorch.org/uploads/default/6d0d5f2696fa6a513e075fcb2bf76d89af915f96" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/6/d/6d0d5f2696fa6a513e075fcb2bf76d89af915f96.png" alt="image" data-base62-sha1="fyIGXNLmxXov1OayhjszsttHiwS" width="690" height="86" data-dominant-color="491517"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1046×131 5.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

ptrblck: <p>Could you post the shape of <code>target</code> as it seems the class indices is out of bounds created via <code>torch.argmax(action).item()</code>?</p>

Bocksa: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="199986">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>torch.argmax(action).item()</p>
</blockquote>
</aside>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/e/be7275eff98d9e9434880ed633e9eefa21fd87b5.png" data-download-href="https://discuss.pytorch.org/uploads/default/be7275eff98d9e9434880ed633e9eefa21fd87b5" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/e/be7275eff98d9e9434880ed633e9eefa21fd87b5_2_690x236.png" alt="image" data-base62-sha1="raLWUUtX0BTyI4Mrbi12IOKUqLb" width="690" height="236" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/e/be7275eff98d9e9434880ed633e9eefa21fd87b5_2_690x236.png, https://discuss.pytorch.org/uploads/default/optimized/3X/b/e/be7275eff98d9e9434880ed633e9eefa21fd87b5_2_1035x354.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/b/e/be7275eff98d9e9434880ed633e9eefa21fd87b5.png 2x" data-dominant-color="242523"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1219×417 25.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

Bocksa: <p>Changed batch size to 3 and it seemed to have worked, unsure why.</p>