add mlp to device

alexhernandezgarcia · Sep 23, 2024 · 9395e61 · 9395e61
1 parent 2bf438a
commit 9395e61
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/gflownet/policy/mlp.py b/gflownet/policy/mlp.py
@@ -37,7 +37,6 @@ def make_model(self, activation: nn.Module = nn.LeakyReLU()):
         is_model : bool
             True because an MLP is a model.
         """
-        activation.to(self.device)
 
         if self.shared_weights == True and self.base is not None:
             mlp = nn.Sequential(
@@ -66,7 +65,7 @@ def make_model(self, activation: nn.Module = nn.LeakyReLU()):
                     + self.tail
                 )
             )
-            return mlp, True
+            return mlp.to(self.device), True
         else:
             raise ValueError(
                 "Base Model must be provided when shared_weights is set to True"