mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
c797fbc4e1
Partially addresses #123062 Ran lintrunner on: - `test/cpp_api_parity` - `test/cpp_extensions` - `test/create_dummy_torchscript_model.py` - `test/custom_backend` - `test/custom_operator` Pull Request resolved: https://github.com/pytorch/pytorch/pull/123518 Approved by: https://github.com/huydhn
37 lines
892 B
Python
37 lines
892 B
Python
# Usage: python create_dummy_model.py <name_of_the_file>
|
|
import sys
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
|
|
class NeuralNetwork(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.flatten = nn.Flatten()
|
|
self.linear_relu_stack = nn.Sequential(
|
|
nn.Linear(28 * 28, 512),
|
|
nn.ReLU(),
|
|
nn.Linear(512, 512),
|
|
nn.ReLU(),
|
|
nn.Linear(512, 10),
|
|
)
|
|
|
|
def forward(self, x):
|
|
x = self.flatten(x)
|
|
logits = self.linear_relu_stack(x)
|
|
return logits
|
|
|
|
|
|
if __name__ == "__main__":
|
|
jit_module = torch.jit.script(NeuralNetwork())
|
|
torch.jit.save(jit_module, sys.argv[1])
|
|
orig_module = nn.Sequential(
|
|
nn.Linear(28 * 28, 512),
|
|
nn.ReLU(),
|
|
nn.Linear(512, 512),
|
|
nn.ReLU(),
|
|
nn.Linear(512, 10),
|
|
)
|
|
torch.save(orig_module, sys.argv[1] + ".orig")
|