add raw_logits in categorical to() method (#131)

recursionpharma · May 6, 2024 · 02eaed3 · 02eaed3
1 parent 286c301
commit 02eaed3
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/gflownet/envs/graph_building_env.py b/src/gflownet/envs/graph_building_env.py
@@ -627,6 +627,7 @@ def detach(self):
 
     def to(self, device):
         self.dev = device
+        self.raw_logits = [i.to(device) for i in self.raw_logits]
         self._masked_logits = [i.to(device) for i in self._masked_logits]
         self.batch = [i.to(device) for i in self.batch]
         self.slice = [i.to(device) for i in self.slice]