randomwalks.py doesn't work on v0.3 or current main

[dpaleka]

🐛 Describe the bug

On both v0.3 and https://github.com/CarperAI/trlx/commit/ff0d0776ce9189c7e0ebc954dd14bbca1136a450, following the instructions from README.md and running

wandb disable && python examples/randomwalks.py

produces the following error:

Traceback (most recent call last):
  File "/home/dpaleka/code/trlx/examples/randomwalks.py", line 103, in <module>
    trlx.train(
  File "/home/dpaleka/code/trlx/trlx/trlx.py", line 95, in train
    model.learn()
  File "/home/dpaleka/code/trlx/trlx/model/accelerate_base_model.py", line 240, in learn
    results = self.evaluate()
  File "/home/dpaleka/code/trlx/trlx/model/accelerate_base_model.py", line 160, in evaluate
    samples = self.generate(prompts)
  File "/home/dpaleka/code/trlx/trlx/model/accelerate_base_model.py", line 133, in generate
    return self.accelerator.unwrap_model(self.model).generate(
  File "/home/dpaleka/code/trlx/trlx/model/nn/ilql_models.py", line 306, in generate
    logits[torch.where(logit_mask[input_ids[:, -1].squeeze()])] = -np.inf
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

Which trlX version are you using?

trlx==0.3

Additional system and package information

No response

[LouisCastricato]

Resolved #105