Skip to content
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions sendnn_inference/v1/worker/spyre_model_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -328,11 +328,11 @@ def load_model(self) -> None:
self._model = self._model.base_model
else:
self._model = AutoModel.from_pretrained(
self.model_config.model, dtype=torch.float16
self.model_config.model, dtype=torch.float16, attn_implementation="eager"
)
elif task == "classify":
class_model = AutoModelForSequenceClassification.from_pretrained(
self.model_config.model
self.model_config.model, dtype=torch.float16, attn_implementation="eager"
)
if hasattr(class_model, "bert"):
self._model = class_model.bert
Expand Down
Loading