compute_train_stats: Fix logits passed in as proba

[shwang]

Fixes an ~"negative probability out-of-bounds" error I ran into when I was training.

See: https://pytorch.org/docs/stable/distributions.html#bernoulli for func signature

[codecov]

Codecov Report

Merging #273 (bcae317) into master (9ddd6fa) will not change coverage.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master     #273   +/-   ##
=======================================
  Coverage   90.01%   90.01%           
=======================================
  Files          74       74           
  Lines        5146     5146           
=======================================
  Hits         4632     4632           
  Misses        514      514           

Impacted Files	Coverage Δ
src/imitation/rewards/common.py	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 9ddd6fa...bcae317. Read the comment docs.

[AdamGleave]

LGTM.

This seems like a badly designed function: Bernoulli(probs, logits) is illegal (can only specify one), you have to specify logits via keyword argument. Would have been better if it had been Bernoulli(*, probs, logits) to force people to always use keyword arguments. Oh well.