| 55 | |
| 56 | class LlavaEdgeManager(LLMEdgeManager): |
| 57 | def export(self) -> "LlavaEdgeManager": |
| 58 | dynamic_shape = self._get_dynamic_shape() |
| 59 | # 1. torch.nn.attention.sdpa_kernel([SDPBackend.MATH]) is for bypassing the dynamo error when tracing |
| 60 | # 2. torch.no_grad() is for getting rid of the dropout (not sure why training ops will show up) |
| 61 | with torch.nn.attention.sdpa_kernel([SDPBackend.MATH]), torch.no_grad(): |
| 62 | self.export_program = torch.export.export( |
| 63 | self.model, |
| 64 | self.example_inputs, |
| 65 | dynamic_shapes=dynamic_shape, |
| 66 | strict=False, |
| 67 | ) |
| 68 | self.pre_autograd_graph_module = self.export_program.module() |
| 69 | return self |
| 70 | |
| 71 | |
| 72 | def export_text_model(llava, embeddings, dynamic_shapes): |