Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

quick fix to continue with issue 71 #73

Closed
wants to merge 1 commit into from
Closed

quick fix to continue with issue 71 #73

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 3 additions & 4 deletions tritonbench/operator_loader/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
from typing import Any, Generator, List, Optional

import torch
from torch._dynamo.backends.cudagraphs import cudagraphs_inner
from torch._inductor.utils import gen_gm_and_inputs
from torch._ops import OpOverload
from torch.utils._pytree import tree_map_only
Expand Down Expand Up @@ -85,6 +84,9 @@ def __init__(
), f"AtenOpBenchmark only supports fp16 and fp32, but got {self.dtype}"

def get_input_iter(self) -> Generator:
from torch._dynamo.backends.cudagraphs import cudagraphs_inner
from torch._inductor.compile_fx import compile_fx

inps_gens = [self.huggingface_loader, self.torchbench_loader, self.timm_loader]
for inp_gen in inps_gens:
for inp in inp_gen.get_inputs_for_operator(
Expand All @@ -101,9 +103,6 @@ def get_input_iter(self) -> Generator:
"aten::convolution_backward",
)
if self.device == "cuda":
from torch._inductor.compile_fx import compile_fx


cudagraph_eager = cudagraphs_inner(
gm, gm_args, copy_outputs=False, copy_inputs=False
)
Expand Down
Loading