Pyro.plate error with torch.diag()

[FreeAJust]

Thanks! It helps a lot!