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"))

[((153, 231), 'pytest.mark.filterwarnings', 'pytest.mark.filterwarnings', (['"""ignore:Named tensors .* experimental:UserWarning"""'], {}), "('ignore:Named tensors .* experimental:UserWarning')\n", (179, 231), False, 'import pytest\n'), ((502, 578), 'pytest.mark.filterwarnings', 'pytest.mark.filterwarnings', (['"""ignore:.*ONNX contains stripped .*:UserWarning"""'], {}), "('ignore:.*ONNX contains stripped .*:UserWarning')\n", (528, 578), False, 'import pytest\n'), ((272, 277), 'tests.pytorch_pfn_extras_tests.onnx.test_export_testcase.Net', 'Net', ([], {}), '()\n', (275, 277), False, 'from tests.pytorch_pfn_extras_tests.onnx.test_export_testcase import Net\n'), ((628, 633), 'tests.pytorch_pfn_extras_tests.onnx.test_export_testcase.Net', 'Net', ([], {}), '()\n', (631, 633), False, 'from tests.pytorch_pfn_extras_tests.onnx.test_export_testcase import Net\n'), ((344, 368), 'torch.rand', 'torch.rand', (['(1)', '(1)', '(28)', '(28)'], {}), '(1, 1, 28, 28)\n', (354, 368), False, 'import torch\n'), ((463, 497), 'os.path.join', 'os.path.join', (['outdir', '"""model.onnx"""'], {}), "(outdir, 'model.onnx')\n", (475, 497), False, 'import os\n'), ((709, 733), 'torch.rand', 'torch.rand', (['(1)', '(1)', '(28)', '(28)'], {}), '(1, 1, 28, 28)\n', (719, 733), False, 'import torch\n'), ((882, 916), 'os.path.join', 'os.path.join', (['outdir', '"""model.onnx"""'], {}), "(outdir, 'model.onnx')\n", (894, 916), False, 'import os\n')]