Encounter an CUDA error when set Medusa head

[1649759610]

hi, @ctlllll

I try to use medusa on llama model，and do some medusa head experiments.

when base_model_config. medusa_num_heads in from_pretrained(medusa_model.py) is set to be 2 or 3， an error will be raised as follow；but if set to be 5, it seems to work well, Could you tell its reason for this?