Create new BEST_CONFIGS when loading wasn't possible (#73)

pytorch-labs · Nov 18, 2023 · 2ed7d1a · 2ed7d1a
1 parent 0b599aa
commit 2ed7d1a
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/segment_anything_fast/flash_4.py b/segment_anything_fast/flash_4.py
@@ -292,6 +292,9 @@ def _attention_rel_h_rel_w_kernel_aligned(q, k, v, rel_h_w, sm_scale):
     global BEST_CONFIGS
     if BEST_CONFIGS is None:
         BEST_CONFIGS = _load_best_configs()
+    # Loading must have not been successful. Let's create a new dictionary.
+    if BEST_CONFIGS is None:
+        BEST_CONFIGS = {}
     key = _create_best_configs_key(q, k, v, rel_h_w, o)
     if key not in BEST_CONFIGS:
         print("key ", key, " not found. Running autotune. This might take a while.")