diff options
Diffstat (limited to 'res/PyTorchExamples/examples/split-1/__init__.py')
-rw-r--r-- | res/PyTorchExamples/examples/split-1/__init__.py | 17 |
1 files changed, 17 insertions, 0 deletions
diff --git a/res/PyTorchExamples/examples/split-1/__init__.py b/res/PyTorchExamples/examples/split-1/__init__.py new file mode 100644 index 000000000..33811fd44 --- /dev/null +++ b/res/PyTorchExamples/examples/split-1/__init__.py @@ -0,0 +1,17 @@ +import torch +import torch.nn as nn + + +# model +class net_split(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, input): + return torch.split(input, (1, 3)) + + +_model_ = net_split() + +# dummy input for onnx generation +_dummy_ = torch.randn(4, 2, 3, 3) |