mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 21:49:24 +08:00
042d764872
Summary: To give user a simpler example code, we are getting rid of ExportArgs in favor of example_args and example_kwargs. Test Plan: CI Differential Revision: D59288920 Pull Request resolved: https://github.com/pytorch/pytorch/pull/129982 Approved by: https://github.com/angelayi
18 lines
404 B
Python
18 lines
404 B
Python
# mypy: allow-untyped-defs
|
|
import torch
|
|
|
|
class Dictionary(torch.nn.Module):
|
|
"""
|
|
Dictionary structures are inlined and flattened along tracing.
|
|
"""
|
|
|
|
def forward(self, x, y):
|
|
elements = {}
|
|
elements["x2"] = x * x
|
|
y = y * elements["x2"]
|
|
return {"y": y}
|
|
|
|
example_args = (torch.randn(3, 2), torch.tensor(4))
|
|
tags = {"python.data-structure"}
|
|
model = Dictionary()
|