Skip to content

Commit

Permalink
changed the signature of the torch compile wrarpper function (#3139)
Browse files Browse the repository at this point in the history
### Changes

Updated get_torch_compile_wrapper

### Related tickets

#3134

### Tests
the current test scope
  • Loading branch information
alexsu52 authored Dec 12, 2024
1 parent 3ff9fe1 commit 42d4c9a
Showing 1 changed file with 3 additions and 3 deletions.
6 changes: 3 additions & 3 deletions nncf/torch/dynamic_graph/patch_pytorch.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
import functools
import inspect
from contextlib import contextmanager
from typing import Callable, List, Union
from typing import Callable, List, Optional, Union

import torch
import torch.utils.cpp_extension
Expand Down Expand Up @@ -251,13 +251,13 @@ def get_torch_compile_wrapper():
"""

@functools.wraps(_ORIG_TORCH_COMPILE)
def wrapper(model, *args, **kwargs):
def wrapper(model: Optional[Callable] = None, **kwargs):
from nncf.torch.nncf_network import NNCFNetwork

if isinstance(model, NNCFNetwork):
raise TypeError("At the moment torch.compile() is not supported for models optimized by NNCF.")
with disable_patching():
return _ORIG_TORCH_COMPILE(model, *args, **kwargs)
return _ORIG_TORCH_COMPILE(model, **kwargs)

return wrapper

Expand Down

0 comments on commit 42d4c9a

Please sign in to comment.