Trying to move tensors not parameters to device

Signed-off-by: João Lucas de Sousa Almeida <[email protected]>
IBM · Feb 7, 2024 · a0b8ab5 · a0b8ab5
1 parent 5a9e73f
commit a0b8ab5
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/simulai/activations.py b/simulai/activations.py
@@ -102,13 +102,13 @@ def __init__(self, device:str="cpu") -> None:
         super(Wavelet, self).__init__()
 
         self.device = device
-        self.w1 = torch.nn.Parameter(torch.ones(1), requires_grad=True).to(self.device)
-        self.w2 = torch.nn.Parameter(torch.ones(1), requires_grad=True).to(self.device)
+        self.w1 = torch.nn.Parameter(torch.ones(1).to(self.device), requires_grad=True)
+        self.w2 = torch.nn.Parameter(torch.ones(1).to(self.device), requires_grad=True)
 
     def setup(self, device:str=None) -> None:
 
-        self.w1 = torch.nn.Parameter(torch.ones(1), requires_grad=True).to(device)
-        self.w2 = torch.nn.Parameter(torch.ones(1), requires_grad=True).to(device)
+        self.w1 = torch.nn.Parameter(torch.ones(1).to(self.device), requires_grad=True)
+        self.w2 = torch.nn.Parameter(torch.ones(1).to(self.device), requires_grad=True)
 
     def forward(self, input: torch.Tensor) -> torch.Tensor:
         """Perform the forward pass of the Wavelet activation on the input.