Skip to content

Commit

Permalink
modify perplexity test to use sdpa
Browse files Browse the repository at this point in the history
  • Loading branch information
dan-garvey committed Nov 1, 2024
1 parent 344cda2 commit 5782f8b
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion sharktank/sharktank/evaluate/perplexity_torch.py
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,7 @@ def load_model(self, dataset, tokenizer, tensor_parallelism_size, attention_kern
attention_dtype=self.attention_dtype,
tensor_parallelism_size=tensor_parallelism_size,
)

config.attention_kernel="torch"
if config.tensor_parallelism_size > 1:
dataset.root_theta = shard_theta(dataset.root_theta, config)

Expand Down

0 comments on commit 5782f8b

Please sign in to comment.