Fix: loading DBRX back from saved path (#35728)

* fix dtype as dict for some models + add test

* add comment in tests
This commit is contained in:
Raushan Turganbay
2025-01-28 11:38:45 +01:00
committed by GitHub
parent 3613f568cd
commit b764c20b09
4 changed files with 13 additions and 4 deletions

View File

@@ -331,6 +331,12 @@ class ModelTesterMixin:
with torch.no_grad():
second = model(**self._prepare_for_class(inputs_dict, model_class))[0]
# Save and load second time because `from_pretrained` adds a bunch of new config fields
# so we need to make sure those fields can be loaded back after saving
# Simply init as `model(config)` doesn't add those fields
model.save_pretrained(tmpdirname)
model = model_class.from_pretrained(tmpdirname)
if isinstance(first, tuple) and isinstance(second, tuple):
for tensor1, tensor2 in zip(first, second):
check_save_load(tensor1, tensor2)