reshape after view is wrong

**Describe the bug**
Linear followed by a view leads to incorrect result, tensorflowjs complains shape mismatch [NULL, output_size]. changing 
`    reshape = keras.layers.Reshape([-1], name=tf_name)` to match the input shape seems to fix the problem

**To Reproduce**
__Snippet of your code__

```
x = x.view(self.batch_size, -1)
self.output = nn.Linear(x.shape[1], output_size)
```

