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])
def test_serve():
model = SemanticSegmentation(2)
# TODO: Currently only servable once a preprocess has been attached
model._preprocess = SemanticSegmentationPreprocess()
model.eval()
model.serve()
def test_serve():
model = SemanticSegmentation(2)
model.eval()
model.serve()