Vim
Vim copied to clipboard
Run the test code successfully, share my environment.
Hello every researcher, I have researched this code for 3 day, Finally, I can run the test code successfully. below is the test code. `import torch from mamba_ssm import Mamba
batch, length, dim = 2, 64, 16 print(torch.cuda.is_available()) x = torch.randn(batch, length, dim).to("cuda") model = Mamba( # This module uses roughly 3 * expand * d_model^2 parameters d_model=dim, # Model dimension d_model d_state=16, # SSM state expansion factor d_conv=4, # Local convolution width expand=2, # Block expansion factor ).to("cuda") y = model(x) assert y.shape == x.shape`
Therefore, I would like to share my experiments to help those people are still struggling for the problem of setting environment.
my syetem:
**Ubuntu 22.04
nvidia driver 535 (if you don't know how to upgrade or downgrade your nvidia-driver, you can leave a message here)
CUDA 11.8 (downloaded run file from official website)
python 3.10.13 (installed via conda )
pytorch 2.2.2 (installed via conda)
causal-conv1d-1.1.3.post1( installed via pip install causal-conv1d==1.1.3.post1
)
mamba_ssm(installed via author provided file cd mamba-1p1p1/ pip install -e .
)**