Spaces:

huggingface
/

keras-chatbot-arena

Running on TPU v5e

martin-gorner HF staff commited on 27 days ago

Commit

38f8411

•

1 Parent(s): d20c9b9

bug fix

Files changed (1) hide show

models.py CHANGED Viewed

@@ -60,7 +60,8 @@ def get_default_layout_map(preset_name, device_mesh):
             patch_key = "decoder_block.*attention.*(query|key|value).kernel"
             layout_map.pop(patch_key)
             layout_map[patch_key] = (None, "model", "batch")
-            return layout_map
 def log_applied_layout_map(model):

             patch_key = "decoder_block.*attention.*(query|key|value).kernel"
             layout_map.pop(patch_key)
             layout_map[patch_key] = (None, "model", "batch")
+        return layout_map
 def log_applied_layout_map(model):