Fix idefics.

huggingface · Aug 14, 2024 · b2d3948 · b2d3948
1 parent 5c598cc
commit b2d3948
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/server/text_generation_server/models/idefics.py b/server/text_generation_server/models/idefics.py
@@ -33,6 +33,7 @@ def __init__(
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.quantize = quantize
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():
             device = torch.device(f"cuda:{rank}")

diff --git a/server/text_generation_server/models/idefics_causal_lm.py b/server/text_generation_server/models/idefics_causal_lm.py
@@ -580,6 +580,7 @@ def __init__(
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.quantize = quantize
         from text_generation_server.models.custom_modeling.idefics_modeling import (
             IdeficsForVisionText2Text,
         )