ai21labs

/

Jamba-v0.1

like 1.17k

Follow

AI21 180

I want to load model with bf16, and evaluate it on some dataset.
But when I use zero3 to distrubute model weight into different GPUs, there is a error showing as below:

modeling_jamba.py", line 913, in cuda_kernels_forward
    conv_weights = self.conv1d.weight.view(self.conv1d.weight.size(0), self.conv1d.weight.size(2))
IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 2)

So can you tell me how to evaluate Jamba in bf16 using accelrate ?