gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[TESTS] Added bfloat16 tests (#430)

Browse Source
This commit is contained in:
daadaada
2022-01-14 15:38:32 +08:00
committed by GitHub
parent 4c94359199
commit 2a944ded53
6 changed files with 22 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/triton/code_gen.py
View File

@@ -838,7 +838,7 @@ class Autotuner:
# prune configs
pruned_configs = self.configs
if self.prune_num_stages_by:
pruned_configs = self.prune_num_stages_by(self.configs)
pruned_configs = self.prune_num_stages_by(self.configs, self.nargs)
if self.perf_model:
top_k = self.configs_top_k
if isinstance(top_k, float) and top_k <= 1.0: