Skip to content

GitLab

Unverified Commit a97a5a9e authored by Hailey Schoelkopf's avatar Hailey Schoelkopf Committed by GitHub
Browse files

make optimum imports conditional

parent 84f1a131
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 2 deletions
lm_eval/models/gpt2.py
View file @ a97a5a9e
...@@ -2,8 +2,6 @@ import torch ...@@ -2,8 +2,6 @@ import torch
import transformers import transformers
from typing import Optional, Union from typing import Optional, Union
from lm_eval.base import BaseLM from lm_eval.base import BaseLM
import optimum
from optimum.intel.openvino import OVModelForCausalLM
def _get_dtype(dtype: Union[str, torch.dtype]) -> torch.dtype: def _get_dtype(dtype: Union[str, torch.dtype]) -> torch.dtype:
...@@ -193,6 +191,9 @@ class OPTIMUMLM(BaseLM): ...@@ -193,6 +191,9 @@ class OPTIMUMLM(BaseLM):
): ):
super().__init__() super().__init__()
import optimum
from optimum.intel.openvino import OVModelForCausalLM
assert isinstance(device, str) assert isinstance(device, str)
assert isinstance(pretrained, str) assert isinstance(pretrained, str)
assert isinstance(batch_size, (int,str)) assert isinstance(batch_size, (int,str))
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment