Fix bug in sampler log_prob dim

[alexnikulkov]

Summary: Fixing a bug introduced in D41062175

Differential Revision: D41164406

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D41164406

[codecov-commenter]

Codecov Report

Base: 87.63% // Head: 69.54% // Decreases project coverage by -18.09% :warning:

Coverage data is based on head (eff2128) compared to base (4ea529e). Patch coverage: 0.00% of modified lines in pull request are covered.

Additional details and impacted files

```diff @@ Coverage Diff @@ ## main #692 +/- ## =========================================== - Coverage 87.63% 69.54% -18.10% =========================================== Files 365 364 -1 Lines 23678 23622 -56 Branches 44 44 =========================================== - Hits 20751 16427 -4324 - Misses 2901 7169 +4268 Partials 26 26 ``` | [Impacted Files](https://codecov.io/gh/facebookresearch/ReAgent/pull/692?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=facebookresearch) | Coverage Δ | | |---|---|---| | [reagent/gym/policies/samplers/discrete\_sampler.py](https://codecov.io/gh/facebookresearch/ReAgent/pull/692/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=facebookresearch#diff-cmVhZ2VudC9neW0vcG9saWNpZXMvc2FtcGxlcnMvZGlzY3JldGVfc2FtcGxlci5weQ==) | `83.15% <0.00%> (ø)` | | | [reagent/test/base/test\_types.py](https://codecov.io/gh/facebookresearch/ReAgent/pull/692/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=facebookresearch#diff-cmVhZ2VudC90ZXN0L2Jhc2UvdGVzdF90eXBlcy5weQ==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: | | [reagent/test/base/test\_utils.py](https://codecov.io/gh/facebookresearch/ReAgent/pull/692/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=facebookresearch#diff-cmVhZ2VudC90ZXN0L2Jhc2UvdGVzdF91dGlscy5weQ==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: | | [reagent/test/core/test\_utils.py](https://codecov.io/gh/facebookresearch/ReAgent/pull/692/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=facebookresearch#diff-cmVhZ2VudC90ZXN0L2NvcmUvdGVzdF91dGlscy5weQ==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: | | [reagent/test/models/test\_bcq.py](https://codecov.io/gh/facebookresearch/ReAgent/pull/692/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=facebookresearch#diff-cmVhZ2VudC90ZXN0L21vZGVscy90ZXN0X2JjcS5weQ==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: | | [reagent/test/models/test\_dqn.py](https://codecov.io/gh/facebookresearch/ReAgent/pull/692/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=facebookresearch#diff-cmVhZ2VudC90ZXN0L21vZGVscy90ZXN0X2Rxbi5weQ==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: | | [reagent/test/models/test\_actor.py](https://codecov.io/gh/facebookresearch/ReAgent/pull/692/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=facebookresearch#diff-cmVhZ2VudC90ZXN0L21vZGVscy90ZXN0X2FjdG9yLnB5) | `0.00% <0.00%> (-100.00%)` | :arrow_down: | | [reagent/test/models/test\_critic.py](https://codecov.io/gh/facebookresearch/ReAgent/pull/692/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=facebookresearch#diff-cmVhZ2VudC90ZXN0L21vZGVscy90ZXN0X2NyaXRpYy5weQ==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: | | [reagent/test/core/aggregators\_test.py](https://codecov.io/gh/facebookresearch/ReAgent/pull/692/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=facebookresearch#diff-cmVhZ2VudC90ZXN0L2NvcmUvYWdncmVnYXRvcnNfdGVzdC5weQ==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: | | [reagent/test/base/test\_tensorboardX.py](https://codecov.io/gh/facebookresearch/ReAgent/pull/692/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=facebookresearch#diff-cmVhZ2VudC90ZXN0L2Jhc2UvdGVzdF90ZW5zb3Jib2FyZFgucHk=) | `0.00% <0.00%> (-100.00%)` | :arrow_down: | | ... and [85 more](https://codecov.io/gh/facebookresearch/ReAgent/pull/692/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=facebookresearch) | | Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=facebookresearch). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=facebookresearch)

:umbrella: View full report at Codecov.
:loudspeaker: Do you have feedback about the report comment? Let us know in this issue.