Inference problem

[Artem535]

Hi! 👋 I'm trying to use Unsloth for inference with the model "unsloth/Qwen3-8B-unsloth-bnb-4bit", but I encountered this error:

/python3.12/site-packages/unsloth/kernels/cross_entropy_loss.py", line 400, in fast_cross_entropy_loss
    assert(labels.shape == (batch, seq_len)

It looks like a shape mismatch, but I'm unsure where it’s coming from since I’m only doing inference (not training).

Here’s the code snippet I’m using to initialize the model:

def _get_model(self):
    torch.set_float32_matmul_precision("high")
    model, _  = FastLanguageModel.from_pretrained(
        self.config.llm_model.llm_url,
        device_map=self.config.environment.device_type,
        dtype=torch.bfloat16,
        load_in_4bit=self.config.llm_model.load_in_4bit,
    )
    model = FastLanguageModel.for_inference(model)
    return model

Any idea how I can resolve this error or what could be causing it?