Solved by vdw in post #8 
              
              
                I think I found the reason show it works despite the wrong flattening. You’re last two lines in your forward() method are: 
def forward(self, x, hs):
    ...
    out = self.fc(out)
    return out[-1], hs

The out[-1] resolves the “artificial” batch of 12 values so you have only 1 output value. If yo…