| 439 | image = inputs['flattened_patches'].to(args.device, dtype) |
| 440 | |
| 441 | class pix2structVisionWrapper(torch.nn.Module): |
| 442 | |
| 443 | def __init__(self, encoder): |
| 444 | super().__init__() |
| 445 | self.encoder = encoder |
| 446 | |
| 447 | def forward(self, image): |
| 448 | attention_mask = (image.abs().sum(dim=-1) != 0) |
| 449 | vision_x = self.encoder.embeddings(image) |
| 450 | img_features = self.encoder.encoder(vision_x, |
| 451 | attention_mask=attention_mask) |
| 452 | img_features = self.encoder.layernorm(img_features[0]) |
| 453 | return img_features |
| 454 | |
| 455 | model = Pix2StructForConditionalGeneration.from_pretrained(args.model_path, |
| 456 | dtype=dtype) |
no outgoing calls
no test coverage detected