Make assertions only if actually chunking forward
This moves the assertion on checking input dimensions into a block that will only be called if the function is actually going to do chunking forward. This is often not the case at inference time and PyTorch tracing a model with this assertion in it leads to a tracing warning.

```
TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  input_tensor.shape[chunk_dim] == tensor_shape for input_tensor in input_tensors
```
Asking @patrickvonplaten for review as it's on reformer<|||||>@LysandreJik This appears when tracing [sentence-transformers/all-MiniLM-L12-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L12-v2), which doesn't seem to be a reformer 🤔 <|||||>Hey @joshdevins, 

I'm happy with the PR! Can you rebase to current master - it should fix the failing tests I think. Also could you maybe replace the `assert` statement by `if ... raise ValueError`? We are trying to move away from assert statements :-) 