From ded5e3e624bdb7e3b579b859c8b6255b47448532 Mon Sep 17 00:00:00 2001 From: Maks Sorokin Date: Fri, 19 Apr 2024 10:36:08 -0400 Subject: [PATCH] change device source --- torchrl/collectors/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchrl/collectors/utils.py b/torchrl/collectors/utils.py index b043445bf94..6888b8b043c 100644 --- a/torchrl/collectors/utils.py +++ b/torchrl/collectors/utils.py @@ -179,7 +179,7 @@ def split_trajectories( torch.ones( out_split.shape, dtype=torch.bool, - device=out_split.get(("next", "done")).device, + device=out_split.device, ), ) if len(out_splits) > 1: