fix pyro global step tracking

[adamgayoso]

The issue here is that PyTorch Lightning uses its wrapper of an optimizer to track steps. It uses post step hooks to increment the global step of the trainer.

Here we provide a shimmy optimizer for it to do this.

[codecov[bot]]

Codecov Report

Merging #1636 (5857305) into master (428d469) will increase coverage by 0.02%. The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master    #1636      +/-   ##
==========================================
+ Coverage   91.24%   91.26%   +0.02%     
==========================================
  Files         116      116              
  Lines        9192     9196       +4     
==========================================
+ Hits         8387     8393       +6     
+ Misses        805      803       -2     

Impacted Files	Coverage Δ
scvi/train/_trainingplans.py	95.66% <100.00%> (+0.03%)	:arrow_up:
scvi/data/_download.py	87.17% <0.00%> (+0.33%)	:arrow_up:
scvi/__init__.py	100.00% <0.00%> (+13.33%)	:arrow_up:

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[vitkl]

Thanks for solving this issue! I also get this issue with non-cell2location scvi-tools model. So would be great if this fix is approved and merged.