Skip to content

Commit

Permalink
[BugFix] Do not unlock td if it's not locked in TDParams (for compile…
Browse files Browse the repository at this point in the history
… compat)

ghstack-source-id: 9b6923f9c219e12af5560c97c1c6c58ed7870a8a
Pull Request resolved: #1125
  • Loading branch information
vmoens committed Dec 3, 2024
1 parent dd88950 commit 4949efa
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion tensordict/nn/params.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import re
import weakref
from concurrent.futures import Future, ThreadPoolExecutor
from contextlib import nullcontext
from copy import copy
from functools import wraps
from typing import Any, Callable, Dict, Iterator, List, OrderedDict, Sequence, Type
Expand Down Expand Up @@ -171,7 +172,11 @@ def new_func(_self, *args, **kwargs):
if _self.is_locked:
# if the root (TensorDictParams) is locked, we still want to raise an exception
raise RuntimeError(_LOCK_ERROR)
with _self._param_td.unlock_():
with (
_self._param_td.unlock_()
if _self._param_td.is_locked
else nullcontext()
):
meth = getattr(_self._param_td, name)
out = meth(*args, **kwargs)
_self._reset_params()
Expand Down

0 comments on commit 4949efa

Please sign in to comment.