-
Mayank Mittal authored
# Description When having agent-environment interaction, it is important to ensure that PyTorch constructs no computational graph. Otherwise, it will keep allocating memory for gradients and result in out-of-memory error. This MR wraps all scripts with `torch.inference_mode()` to prevent this issue. ## Type of change - Bug fix (non-breaking change which fixes an issue) ## Checklist - [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with `./orbit.sh --format` - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] I have updated the changelog and the corresponding version in the extension's `config/extension.toml` file