Skip to content

Commit 3fe17b5

Browse files
committed
Change segformer test to cuda
1 parent 06aafef commit 3fe17b5

1 file changed

Lines changed: 2 additions & 2 deletions

File tree

mlir/benchmarks/torch/model_zoo/segformer_test.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,12 +51,12 @@ def test_backend():
5151
example_input = torch.randn(1, 3, 512, 512)
5252

5353
start = time.perf_counter()
54-
program = torch.compile(model, backend="docc", options={"target": "none", "category": "server"})
54+
program = torch.compile(model, backend="docc", options={"target": "cuda", "category": "server"})
5555
end = time.perf_counter()
5656
print(f"compilation time: {(end - start) * 1000:.2f} ms")
5757

5858
start = time.perf_counter()
59-
ref_program = torch.compile(model, backend="docc", options={"target": "none", "category": "server"})
59+
ref_program = torch.compile(model, backend="docc", options={"target": "cuda", "category": "server"})
6060
end = time.perf_counter()
6161
print(f"ref compilation time: {(end - start) * 1000:.2f} ms")
6262

0 commit comments

Comments
 (0)