You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument mat2 in method wrapper_CUDA_mm
import torch
from zeta.nn import MambaBlock
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
block = MambaBlock(dim=64, depth=1)
x = torch.randn(1, 10, 64).to(device)
y = block(x).to(device)
print(y.shape)
The text was updated successfully, but these errors were encountered:
The same bug:
File "/home/miniconda3/lib/python3.11/site-packages/zeta/nn/modules/simple_mamba.py", line 205, in selective_scan
x = deltaA[:, :, i] * x + deltaB_u[:, :, i]
~~~~~~~~~~~~~~~~^~~
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument mat2 in method wrapper_CUDA_mm
import torch
from zeta.nn import MambaBlock
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
block = MambaBlock(dim=64, depth=1)
x = torch.randn(1, 10, 64).to(device)
y = block(x).to(device)
print(y.shape)
The text was updated successfully, but these errors were encountered: