fix tests

pytorch · Jan 15, 2024 · 121aa73 · 121aa73
1 parent b632be9
commit 121aa73
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/test/test_rb.py b/test/test_rb.py
@@ -1497,8 +1497,8 @@ def test_max_value_writer_reduce(self, size, device, reduction):
  writer=TensorDictMaxValueWriter(rank_key="key", reduction=reduction),
  )
 
- key = torch.rand(batch_size, *size)
- obs = torch.rand(batch_size, *size)
+ key = torch.rand(batch_size, *size, device=device)
+ obs = torch.rand(batch_size, *size, device=device)
  td = TensorDict(
  {"key": key, "obs": obs},
  batch_size=batch_size,