Fix idefics.

This commit is contained in:
Nicolas Patry 2024-08-14 11:14:11 +02:00
parent 5c598cc7ed
commit b2d3948ccf
No known key found for this signature in database
GPG Key ID: 64AF4752B2967863
2 changed files with 2 additions and 0 deletions

View File

@ -33,6 +33,7 @@ class IDEFICSSharded(IdeficsCausalLM):
dtype: Optional[torch.dtype] = None, dtype: Optional[torch.dtype] = None,
trust_remote_code: bool = False, trust_remote_code: bool = False,
): ):
self.quantize = quantize
self.process_group, rank, world_size = initialize_torch_distributed() self.process_group, rank, world_size = initialize_torch_distributed()
if torch.cuda.is_available(): if torch.cuda.is_available():
device = torch.device(f"cuda:{rank}") device = torch.device(f"cuda:{rank}")

View File

@ -580,6 +580,7 @@ class IdeficsCausalLM(Model):
dtype: Optional[torch.dtype] = None, dtype: Optional[torch.dtype] = None,
trust_remote_code: bool = False, trust_remote_code: bool = False,
): ):
self.quantize = quantize
from text_generation_server.models.custom_modeling.idefics_modeling import ( from text_generation_server.models.custom_modeling.idefics_modeling import (
IdeficsForVisionText2Text, IdeficsForVisionText2Text,
) )