Skip to content

Commit

Permalink
fix #56
Browse files Browse the repository at this point in the history
  • Loading branch information
Trinkle23897 committed May 29, 2020
1 parent de556fd commit d2b2fa8
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 1 deletion.
3 changes: 2 additions & 1 deletion test/base/test_batch.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,9 +22,10 @@ def test_batch():

def test_batch_over_batch():
batch = Batch(a=[3, 4, 5], b=[4, 5, 6])
batch2 = Batch(b=batch, c=[6, 7, 8])
batch2 = Batch(c=[6, 7, 8], b=batch)
batch2.b.b[-1] = 0
print(batch2)
assert batch2.values()[-1] == batch2.c
assert batch2[-1].b.b == 0


Expand Down
4 changes: 4 additions & 0 deletions tianshou/data/batch.py
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,10 @@ def keys(self) -> List[str]:
return sorted([
i for i in self.__dict__ if i[0] != '_'] + list(self._meta))

def values(self) -> List[Any]:
"""Return self.values()."""
return [self[k] for k in self.keys()]

def get(self, k: str, d: Optional[Any] = None) -> Union['Batch', Any]:
"""Return self[k] if k in self else d. d defaults to None."""
if k in self.__dict__ or k in self._meta:
Expand Down

0 comments on commit d2b2fa8

Please sign in to comment.