Small bugfix in langevin sampler
Created by: jwnys
Bug fixes to Langevin sampler:
- when providing chunk_size != None, we got a tracing error
- for some models the grad function failed due to a shape (1,)
Created by: jwnys
Bug fixes to Langevin sampler: