import os

import pytest
import torch

import pytorch_pfn_extras.onnx as tou
from tests.pytorch_pfn_extras_tests.onnx.test_export_testcase import Net


@pytest.mark.filterwarnings("ignore:Named tensors .* experimental:UserWarning")
def test_onnx_load_model():
    model = Net()
    outdir = "out/load_model_test"
    tou.export_testcase(model, torch.rand(1, 1, 28, 28), outdir,
                        training=True, do_constant_folding=False)
    tou.load_model(os.path.join(outdir, "model.onnx"))


@pytest.mark.filterwarnings("ignore:.*ONNX contains stripped .*:UserWarning")
def test_stripped_onnx_load_model():
    model = Net()
    outdir = "out/stripped_load_model_test"
    tou.export_testcase(model, torch.rand(1, 1, 28, 28), outdir,
                        strip_large_tensor_data=True, training=True,
                        do_constant_folding=False)
    tou.load_model(os.path.join(outdir, "model.onnx"))

[((10, 1, 10, 79), 'pytest.mark.filterwarnings', 'pytest.mark.filterwarnings', ({(10, 28, 10, 78): '"""ignore:Named tensors .* experimental:UserWarning"""'}, {}), "('ignore:Named tensors .* experimental:UserWarning')", False, 'import pytest\n'), ((19, 1, 19, 77), 'pytest.mark.filterwarnings', 'pytest.mark.filterwarnings', ({(19, 28, 19, 76): '"""ignore:.*ONNX contains stripped .*:UserWarning"""'}, {}), "('ignore:.*ONNX contains stripped .*:UserWarning')", False, 'import pytest\n'), ((12, 12, 12, 17), 'tests.pytorch_pfn_extras_tests.onnx.test_export_testcase.Net', 'Net', ({}, {}), '()', False, 'from tests.pytorch_pfn_extras_tests.onnx.test_export_testcase import Net\n'), ((21, 12, 21, 17), 'tests.pytorch_pfn_extras_tests.onnx.test_export_testcase.Net', 'Net', ({}, {}), '()', False, 'from tests.pytorch_pfn_extras_tests.onnx.test_export_testcase import Net\n'), ((14, 31, 14, 55), 'torch.rand', 'torch.rand', ({(14, 42, 14, 43): '(1)', (14, 45, 14, 46): '(1)', (14, 48, 14, 50): '(28)', (14, 52, 14, 54): '(28)'}, {}), '(1, 1, 28, 28)', False, 'import torch\n'), ((16, 19, 16, 53), 'os.path.join', 'os.path.join', ({(16, 32, 16, 38): 'outdir', (16, 40, 16, 52): '"""model.onnx"""'}, {}), "(outdir, 'model.onnx')", False, 'import os\n'), ((23, 31, 23, 55), 'torch.rand', 'torch.rand', ({(23, 42, 23, 43): '(1)', (23, 45, 23, 46): '(1)', (23, 48, 23, 50): '(28)', (23, 52, 23, 54): '(28)'}, {}), '(1, 1, 28, 28)', False, 'import torch\n'), ((26, 19, 26, 53), 'os.path.join', 'os.path.join', ({(26, 32, 26, 38): 'outdir', (26, 40, 26, 52): '"""model.onnx"""'}, {}), "(outdir, 'model.onnx')", False, 'import os\n')]