@pytest.mark.skipif((not _IMAGE_TESTING), reason="image libraries aren't installed.")
@pytest.mark.parametrize('jitter, args', [(torch.jit.trace, (torch.rand(1, 3, 32, 32),))])
def test_jit(tmpdir, jitter, args):
    path = os.path.join(tmpdir, 'test.pt')
    model = SemanticSegmentation(2)
    model.eval()
    model = jitter(model, *args)
    torch.jit.save(model, path)
    model = torch.jit.load(path)
    out = model(torch.rand(1, 3, 32, 32))
    assert isinstance(out, torch.Tensor)
    assert (out.shape == torch.Size([1, 2, 32, 32]))
@pytest.mark.skipif((not _IMAGE_TESTING), reason="image libraries aren't installed.")
@pytest.mark.parametrize('jitter, args', [(torch.jit.trace, (torch.rand(1, 3, 32, 32),))])
def test_jit(tmpdir, jitter, args):
    path = os.path.join(tmpdir, 'test.pt')
    model = SemanticSegmentation(2)
    model.eval()
    model = jitter(model, *args)
    torch.jit.save(model, path)
    model = torch.jit.load(path)
    out = model(torch.rand(1, 3, 32, 32))
    assert isinstance(out, torch.Tensor)
    assert (out.shape == torch.Size([1, 2, 32, 32]))