Fixed issue #181

kyegomez · Apr 11, 2024 · 22d1599 · 22d1599
1 parent 26052ff
commit 22d1599
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/zeta/nn/modules/simple_mamba.py b/zeta/nn/modules/simple_mamba.py
@@ -199,7 +199,7 @@ def selective_scan(self, u, delta, A, B, C, D):
         )
 
         # Perform selective scan (see scan_SSM() in The Annotated S4 [2])
-        x = torch.zeros((b, d_in, n))
+        x = torch.zeros((b, d_in, n), device=next(self.parameters()).device)
         ys = []
         for i in range(l):
             x = deltaA[:, :, i] * x + deltaB_u[:, :, i]