INNER CODE UNIT · Python
forward
lean-dojo/TorchLean · torchlean/__init__.py:96
def forward(self, *inputs):
return function(*inputs)
model = Function()
model = copy.deepcopy(model)
examples = tuple(value.detach().clone().requires_grad_(value.requires_grad)
for value in examples)
artifact = adapter.export_torchlean_graph_json(
model, examples, str(directory / "graph.json"), require_torch_export=True)
if execute:
artifact["arguments"] = [adapter._constant(value) for value in examples]
if compare:
with torch.no_grad():
outputs, _ = torch.utils._pytree.tree_flatten(model(*examples))
if not all(isinstance(value, torch.Tensor) for value in outputs):
raise ValueError("sample comparison requires tensor outputs")
artifact["reference_outputs"] = [adapter._constant(value) for value in outputs]
return artifact