Neeraj Pradhan

I am not too familiar with the internals, but if it seems like a relatively straightforward fix, I am happy to take a stab at it myself.

> Proper handling though would need similar multiline tricks used by nested bars, but such trickery would have to remain compatible with nested bars. We are using parallel bars too,...

> Hope you know about dynamic_ncols=True. I do! I meant that it would be nice to be able to squash the terminal screen without having the progress bar output get...

The issue is that somewhere during JIT tracing, we are inserting a tensor (most likely a constant) with `requires_grad=True`, whereas `torch.jit` expects all such tensors to be arguments instead. As...

This is a bug which was fixed by @dme65 last month, that's why the nightly seems to work fine. I'll also comment on the issue.

> Easiest fix is probably just to perform that computation on the CPU. Unfortunately, I don't have a good solution in the interim. A *hacky* one that we have previously...

@akihironitta - The codebase and number of contributors to numpyro core is small enough that this is low priority for us at the moment. When I created the issue, I...

I think #20375 is the reason why this fails - the default on Ubuntu and python 3.8 across all platforms is "spawn" context which causes issues. I can replicate it...

@daydreamt - Please go ahead, you will be assured a thorough and timely review. If you have any questions around the codebase, please let us know.

@skye - I thought I'll circle back on this after our brief discussion at Neurips. Our motivation for setting device count is to be able to use `pmap` for running...