Skip to content

Commit

Permalink
[BugFix] unlock in tensordictmodules tests (pytorch#1417)
Browse files Browse the repository at this point in the history
  • Loading branch information
vmoens authored Jul 27, 2023
1 parent e8765c2 commit fc8fe4c
Showing 1 changed file with 10 additions and 5 deletions.
15 changes: 10 additions & 5 deletions test/test_tensordictmodules.py
Original file line number Diff line number Diff line change
Expand Up @@ -857,7 +857,8 @@ def test_functional(self, safe, spec_type):
assert hasattr(tdmodule, "__delitem__")
assert len(tdmodule) == 3
del tdmodule[2]
del params["module", "2"]
with params.unlock_():
del params["module", "2"]
assert len(tdmodule) == 2

assert hasattr(tdmodule, "__getitem__")
Expand Down Expand Up @@ -939,7 +940,8 @@ def test_functional_probabilistic(self, safe, spec_type):
assert hasattr(tdmodule, "__delitem__")
assert len(tdmodule) == 4
del tdmodule[3]
del params["module", "3"]
with params.unlock_():
del params["module", "3"]
assert len(tdmodule) == 3

assert hasattr(tdmodule, "__getitem__")
Expand Down Expand Up @@ -1014,7 +1016,8 @@ def test_functional_with_buffer(self, safe, spec_type):
assert hasattr(tdmodule, "__delitem__")
assert len(tdmodule) == 3
del tdmodule[2]
del params["module", "2"]
with params.unlock_():
del params["module", "2"]
assert len(tdmodule) == 2

assert hasattr(tdmodule, "__getitem__")
Expand Down Expand Up @@ -1103,7 +1106,8 @@ def test_functional_with_buffer_probabilistic(self, safe, spec_type):
assert hasattr(tdmodule, "__delitem__")
assert len(tdmodule) == 4
del tdmodule[3]
del params["module", "3"]
with params.unlock_():
del params["module", "3"]
assert len(tdmodule) == 3

assert hasattr(tdmodule, "__getitem__")
Expand Down Expand Up @@ -1184,7 +1188,8 @@ def test_vmap(self, safe, spec_type):
assert hasattr(tdmodule, "__delitem__")
assert len(tdmodule) == 3
del tdmodule[2]
del params["module", "2"]
with params.unlock_():
del params["module", "2"]
assert len(tdmodule) == 2

assert hasattr(tdmodule, "__getitem__")
Expand Down

0 comments on commit fc8fe4c

Please sign in to comment.