Skip to content

Commit 8c65aa7

Browse files
committed
[BugFix] Fix device initialization in CrossQLoss.maybe_init_target_entropy
Move device retrieval outside the 'if target_entropy == auto' block to ensure the device is available when registering the target_entropy buffer. ghstack-source-id: 515c3a6 Pull-Request: #3311
1 parent fca5dd8 commit 8c65aa7

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

torchrl/objectives/crossq.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -374,8 +374,8 @@ def maybe_init_target_entropy(self, fault_tolerant=True):
374374
if "_target_entropy" in self._buffers:
375375
return
376376
target_entropy = self._target_entropy
377+
device = next(self.parameters()).device
377378
if target_entropy == "auto":
378-
device = next(self.parameters()).device
379379
action_spec = self.get_action_spec()
380380
if action_spec is None:
381381
if fault_tolerant:

0 commit comments

Comments
 (0)