dummy_input = torch.randn(1, 1, 28, 28) torch.onnx.export(model, dummy_input, "stn_classifier.onnx", verbose=True, opset_version=20)
vesion: torch 2.2.2+cu118 torchsummary 1.5.1 torchvision 0.17.2+cu118