Skip to content

Commit

Permalink
Make the test simpler
Browse files Browse the repository at this point in the history
  • Loading branch information
tibor-reiss authored and ArthurZucker committed Nov 19, 2024
1 parent a83e234 commit b782bfb
Showing 1 changed file with 1 addition and 6 deletions.
7 changes: 1 addition & 6 deletions tests/utils/test_modeling_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -467,16 +467,11 @@ def is_on_meta(model_id, dtype):
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=dtype)
return all(value.device.type == "meta" for value in model.state_dict().values())

result = {}
expected = {}
model_ids = ("fxmarty/tiny-llama-fast-tokenizer", "fxmarty/small-llama-testing")
dtypes = (None, "auto", torch.float16)

for model_id, dtype in itertools.product(model_ids, dtypes):
result[(model_id, dtype)] = is_on_meta(model_id, dtype)
expected[(model_id, dtype)] = True

assert result == expected
self.assertTrue(is_on_meta(model_id, dtype))

def test_model_from_pretrained_torch_dtype(self):
# test that the model can be instantiated with dtype of either
Expand Down

0 comments on commit b782bfb

Please sign in to comment.