hiro/ray
1
0
Fork 0
mirror of https://github.com/vale981/ray synced 2025-03-06 02:21:39 -05:00
Code Issues Projects Releases Packages Wiki Activity Actions

[RLlib] Fix torch None conversion in torch_utils.py::convert_to_torch_tensor. (#26863)

Browse source
This commit is contained in:
Rohan Potdar 2022-07-23 04:54:57 -07:00 • committed by GitHub
parent 69f6b843da
commit 97bcf38ec0
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
rllib/utils/torch_utils.py
View file

@ -146,6 +146,9 @@ def convert_to_torch_tensor(x: TensorStructType, device: Optional[str] = None):
"""
def mapping(item):
if item is None:
# returns None with dtype=np.obj
return np.asarray(item)
# Already torch tensor -> make sure it's on right device.
if torch.is_tensor(item):
return item if device is None else item.to(device)