diff --git a/examples/text-generation/utils.py b/examples/text-generation/utils.py index d93b559c3d..6d2381e6e6 100644 --- a/examples/text-generation/utils.py +++ b/examples/text-generation/utils.py @@ -38,7 +38,6 @@ ) from optimum.habana.utils import ( check_habana_frameworks_version, - check_neural_compressor_min_version, check_optimum_habana_min_version, get_habana_frameworks_version, set_seed, @@ -278,8 +277,9 @@ def setup_model(args, model_dtype, model_kwargs, logger): original_model=org_model, **model_kwargs, ) - if not check_neural_compressor_min_version("3.2"): - model = model.to(model_kwargs["torch_dtype"]) + # TODO: This will be removed in v1.19 Synapse release + # the loaded model should have the same dtype as original_model + model = model.to(model_kwargs["torch_dtype"]) else: if args.assistant_model is not None: assistant_model = AutoModelForCausalLM.from_pretrained( diff --git a/optimum/habana/utils.py b/optimum/habana/utils.py index 3e1f755663..a3aa728e81 100755 --- a/optimum/habana/utils.py +++ b/optimum/habana/utils.py @@ -385,15 +385,6 @@ def check_habana_frameworks_version(req_version): ) -def check_neural_compressor_min_version(req_version): - """ - Checks if the installed version of `neural_compressor` is larger than or equal to `req_version`. - """ - import neural_compressor - - return version.Version(neural_compressor.__version__) >= version.Version(req_version) - - def get_device_name(): """ Returns the name of the current device: Gaudi or Gaudi2.