Vim icon indicating copy to clipboard operation
Vim copied to clipboard

Run the test code successfully, share my environment.

Open poult-lab opened this issue 11 months ago • 7 comments

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 .)**

poult-lab avatar Mar 28 '24 06:03 poult-lab