()
| 581 | |
| 582 | |
| 583 | def test_select(): |
| 584 | obs = torch.randn(100, 10) |
| 585 | act = torch.randn(100, 3) |
| 586 | dataset = tu.get_tensordict({"obs": obs, "act": act}, non_tensor_dict={"2": 2, "1": 1}) |
| 587 | |
| 588 | subset = dataset.select("obs", "2") |
| 589 | |
| 590 | assert torch.all(torch.eq(subset["obs"], dataset["obs"])) |
| 591 | assert subset["2"] == dataset["2"] |
| 592 | assert "act" not in subset.keys() |
| 593 | assert "1" not in subset.keys() |
| 594 | |
| 595 | |
| 596 | def test_dataproto_no_batch(): |