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