huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-10 20:04:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Apply suggestions from code review

Browse Source
This commit is contained in:
Nicolas Patry 2023-09-26 15:07:38 +02:00 committed by GitHub
parent 1053e5d09a
commit bf2b92217f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
server/text_generation_server/models/custom_modeling/idefics_vision.py
View File

@ -91,7 +91,7 @@ class IdeficsVisionEmbeddings(nn.Module):
self.position_embedding = TensorParallelEmbedding(
prefix="model.vision_model.embeddings.position_embedding", weights=weights
)
self.position_ids = torch.arange(self.num_positions).expand((1, -1))
self.position_ids = torch.arange(self.num_positions).expand((1, -1)).to(device=weights.device)
def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
batch_size = pixel_values.shape[0]