AssertionError when using model "google/gemma-2b" with multi-gpus

[tritct]

System Info

Information

• Docker
• The CLI directly

Tasks

• An officially supported command
• My own modifications

Reproduction

I'm trying to run Docker on 2 A16 GPUS using model_id "google/gemma-2b". But after the model downloading step I run into AssertionError like the following.

Expected behavior

When I run with only 1 GPU it can initialize just fine. This issue only happen when I try to use multigpu.