| 736 | args.device, torch.float16).unsqueeze(0) |
| 737 | |
| 738 | class FuyuEncoderWrapper(torch.nn.Module): |
| 739 | |
| 740 | def __init__(self, linear): |
| 741 | super().__init__() |
| 742 | self.linear = linear.to(torch.float16) |
| 743 | |
| 744 | def forward(self, patches): |
| 745 | return self.linear(patches).flatten(0, 1) |
| 746 | |
| 747 | model = FuyuForCausalLM.from_pretrained(args.model_path, |
| 748 | dtype=torch.float16) |