[Bug] sv-dkl.py failed when enable jit trace

pyro 1.1.0 sv-dkl.py running error when enable trace. Was exactly running https://github.com/pyro-ppl/pyro/blob/dev/examples/contrib/gp/sv-dkl.py with only --jit set to True on line https://github.com/pyro-ppl/pyro/blob/7692a503b64dc8200049e1c24a775d2170ba8039/examples/contrib/gp/sv-dkl.py#L186 , default setting can successfully be running.

Running in conda (4.7.12) environment with py36-pyro-tutorial.yml set as:

name: py36-pyro-tutorial
dependencies:
  # Core
  - python = 3.6
  - numpy
  - pandas
  - seaborn
  - jupyter
  - torchvision
  - pillow = 6.2.1

  # Language tools
  - pip

  - pip:
    - pyro-ppl == 1.1.0 # 1.0 requires pytorch and torch to be higher version
    - pyro-ppl[extras]

Guidelines

NOTE: Issues are for bugs and feature requests only. If you have a question about using Pyro or general modeling questions, please post it on the forum.

If you would like to address any minor bugs in the documentation or source, please feel free to contribute a Pull Request without creating an issue first.

Please tag the issue appropriately in the title e.g. [bug], [feature request], [discussion], etc.

Please provide the following details:

Issue Description

Provide a brief description of the issue.

Environment

For any bugs, please provide the following:

• OS and python version.
• PyTorch version, or if relevant, output of pip freeze.
• Pyro version: output of python -c 'import pyro; print pyro.__version__'

python 3.6, pyro 1.1.0, conda 4.7.12

Code Snippet

Provide any relevant code snippets and commands run to replicate the issue.

/home/{user}/miniconda3/envs/py36-pyro-tutorial/lib/python3.6/site-packages/pyro/contrib/gp/kernels/isotropic.py:45: 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!
  if X.size(1) != Z.size(1):
/home/{user}/miniconda3/envs/py36-pyro-tutorial/lib/python3.6/site-packages/pyro/contrib/gp/models/vsgp.py:106: TracerWarning: Converting a tensor to a Python index 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!
  Kuu.view(-1)[::M + 1] += self.jitter  # add jitter to the diagonal
/home/{user}/miniconda3/envs/py36-pyro-tutorial/lib/python3.6/site-packages/pyro/distributions/util.py:239: 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!
  eye.view(-1)[:min(m, n) * n:n + 1] = 1
/home/{user}/miniconda3/envs/py36-pyro-tutorial/lib/python3.6/site-packages/pyro/distributions/util.py:239: TracerWarning: Converting a tensor to a Python index 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!
  eye.view(-1)[:min(m, n) * n:n + 1] = 1
/home/{user}/miniconda3/envs/py36-pyro-tutorial/lib/python3.6/site-packages/pyro/contrib/gp/likelihoods/multi_class.py:56: 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!
  if f_swap.size(-1) != self.num_classes:
Traceback (most recent call last):
  File "sv-dkl.py", line 195, in <module>
    main(args)
  File "sv-dkl.py", line 157, in main
    train(args, train_loader, gpmodule, optimizer, loss_fn, epoch)
  File "sv-dkl.py", line 68, in train
    loss = loss_fn(gpmodule.model, gpmodule.guide)
  File "/home/{user}/miniconda3/envs/py36-pyro-tutorial/lib/python3.6/site-packages/pyro/infer/trace_mean_field_elbo.py", line 160, in differentiable_loss
    return self._differentiable_loss(*args, **kwargs)
  File "/home/{user}/miniconda3/envs/py36-pyro-tutorial/lib/python3.6/site-packages/pyro/ops/jit.py", line 92, in __call__
    self.compiled[key] = torch.jit.trace(compiled, params_and_args, **self.jit_options)
  File "/home/{user}/miniconda3/envs/py36-pyro-tutorial/lib/python3.6/site-packages/torch/jit/__init__.py", line 883, in trace
    _force_outplace)
  File "/home/{user}/miniconda3/envs/py36-pyro-tutorial/lib/python3.6/site-packages/pyro/ops/jit.py", line 85, in compiled
    return poutine.replay(self.fn, params=constrained_params)(*args, **kwargs)
  File "/home/{user}/miniconda3/envs/py36-pyro-tutorial/lib/python3.6/site-packages/pyro/poutine/messenger.py", line 8, in _context_wrap
    return fn(*args, **kwargs)
  File "/home/{user}/miniconda3/envs/py36-pyro-tutorial/lib/python3.6/site-packages/pyro/infer/trace_mean_field_elbo.py", line 154, in differentiable_loss
    _, loss_particle = self._differentiable_loss_particle(model_trace, guide_trace)
  File "/home/{user}/miniconda3/envs/py36-pyro-tutorial/lib/python3.6/site-packages/pyro/infer/trace_mean_field_elbo.py", line 99, in _differentiable_loss_particle
    kl_qp = kl_divergence(guide_site["fn"], model_site["fn"])
  File "/home/{user}/miniconda3/envs/py36-pyro-tutorial/lib/python3.6/site-packages/torch/distributions/kl.py", line 165, in kl_divergence
    return fun(p, q)
  File "/home/{user}/miniconda3/envs/py36-pyro-tutorial/lib/python3.6/site-packages/pyro/distributions/kl.py", line 23, in _kl_independent_independent
    kl = kl_divergence(p, q)
  File "/home/{user}/miniconda3/envs/py36-pyro-tutorial/lib/python3.6/site-packages/torch/distributions/kl.py", line 165, in kl_divergence
    return fun(p, q)
  File "/home/{user}/miniconda3/envs/py36-pyro-tutorial/lib/python3.6/site-packages/torch/distributions/kl.py", line 385, in _kl_multivariatenormal_multivariatenormal
    p._unbroadcasted_scale_tril.shape[:-2])