Extend BaichuanForCausalLM to support seq_lens

**Is your feature request related to a problem? Please describe.**
When passing `seq_lens`, currently, I'm getting the following error
```
TypeError: BaichuanForCausalLM.forward() got an unexpected keyword argument 'seq_lens'
```

**Describe the solution you'd like**
Support for seq_lens :) 

**Describe alternatives you've considered**
N/A

**Additional context**
N/A