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