From 2c50d5d09239167c32fdae956d031dcac1e41842 Mon Sep 17 00:00:00 2001 From: Henk Date: Sat, 15 Jul 2023 14:14:06 +0200 Subject: [PATCH] Don't ruin breakmodel --- modeling/inference_models/hf_torch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modeling/inference_models/hf_torch.py b/modeling/inference_models/hf_torch.py index fb9fe39e..1b411c95 100644 --- a/modeling/inference_models/hf_torch.py +++ b/modeling/inference_models/hf_torch.py @@ -412,7 +412,7 @@ class HFTorchInferenceModel(HFInferenceModel): **tf_kwargs, ) - if not self.lazy_load: + if not self.lazy_load and not self.breakmodel: # We need to move the model to the desired device if (not self.usegpu) or torch.cuda.device_count() <= 0: model = model.to("cpu")