I ended up adding these two statements after my linear layer to get final size of output to be 32x4
out = out.view(batch_size, -1, self.output_size)
out = out[:, -1]

Thanks for your help and clarification.