Skip to content

Commit

Permalink
fix
Browse files Browse the repository at this point in the history
  • Loading branch information
bnellnm committed Oct 5, 2024
1 parent 45707e2 commit 52a422c
Showing 1 changed file with 10 additions and 5 deletions.
15 changes: 10 additions & 5 deletions vllm/_custom_ops.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import contextlib
import functools
from typing import List, Optional, Tuple, Union
from typing import TYPE_CHECKING, List, Optional, Tuple, Union

import torch
import torch.library
Expand All @@ -26,10 +26,15 @@
import vllm._moe_C # noqa: F401
supports_moe_ops = True

try:
import torch.library.register_fake
except ImportError:
from torch.library import impl_abstract as register_fake
if TYPE_CHECKING:

def register_fake(fn):
return lambda name: fn
else:
try:
from torch.library import register_fake
except ImportError:
from torch.library import impl_abstract as register_fake


def hint_on_error(fn):
Expand Down

0 comments on commit 52a422c

Please sign in to comment.