#2355 and #2612 conflict

**Describe the bug**

During the merge of #2355, it added code that interacts with #2612 in [2ba8f0b90a6647b869c56350de8cf77798417467](https://github.com/NVIDIA-NeMo/RL/commit/2ba8f0b90a6647b869c56350de8cf77798417467). This, as well as [2ed5993d05d57b0badc0643aa8e564b4c1f63ef2](https://github.com/tdene/RL/commit/2ed5993d05d57b0badc0643aa8e564b4c1f63ef2), must be reconciled after both #2355 and #2612 land.

**Steps/Code to reproduce bug**

Please list *minimal* steps or code snippet for us to be able to reproduce the bug.

A helpful guide on on how to craft a minimal bug report http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports.

**Expected behavior**

A clear and concise description of what you expected to happen.

**Additional context**

Add any other context about the problem here.