I am using svi.step() as loss and found the script stop at the first iteration.
I tried unwrap the svi.step() by defining a function _svi_step() and found the crash happened at
loss = loss.loss_and_grads(model, guide). And again, there is no warning or error log.
The loss is defined as pyro.infer.Trace_ELBO().
Do you have any idea what the issue would be?