The activation parameters were not really being properly device placed.

Signed-off-by: João Lucas de Sousa Almeida <[email protected]>

simulai/models/_pytorch_models/_transformer.py

@@ -3,7 +3,7 @@
 import torch
 from typing import Union, Tuple
 
-from simulai.templates import NetworkTemplate, as_tensor
+from simulai.templates import NetworkTemplate, as_tensor, guarantee_device
 from simulai.regression import DenseNetwork, Linear
 
 

simulai/regression/_pytorch/_dense.py

@@ -200,7 +200,7 @@ def __init__(
 
         """
 
-        super(DenseNetwork, self).__init__()
+        super(DenseNetwork, self).__init__(**kwargs)
 
         assert layers_units, "Please, set a list of units for each layer"
 

simulai/templates/_pytorch_network.py

@@ -55,8 +55,9 @@ def __init__(self, name: str = None, devices: str = None) -> None:
 
         self.shapes_dict = None
         self.device_type = devices
+        self.device = self._set_device(devices=devices)
 
-        if self.device_type:
+        if self.device_type != "cpu":
             self.to_wrap = self._to_explicit_device
         else:
             self.to_wrap = self._to_bypass
@@ -148,7 +149,7 @@ def _get_operation(
                         if torch.nn.Module in res_.__mro__:
                             res = res_
                             print(f"Module {operation} found in {engine}.")
-                            return res()
+                            return res(**kwargs)
                         else:
                             print(f"Module {operation} not found in {engine}.")
                     else:
@@ -175,7 +176,7 @@ def _setup_activations(
             if isinstance(activation_op, simulact.TrainableActivation):
 
                 activations_list = [self._get_operation(operation=activation,
-                                                        is_activation=True, device=self.device_type)
+                                                        is_activation=True, device=self.device)
                                     for i in range(n_layers - 1)]
 
             else: