Simple fix on extended configs

pytorch · Mar 22, 2024 · 38bebdc · 38bebdc
1 parent 9d64293
commit 38bebdc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchbenchmark/util/framework/timm/extended_configs.py b/torchbenchmark/util/framework/timm/extended_configs.py
@@ -19,7 +19,7 @@ def is_extended_timm_models(model_name: str) -> bool:
     return model_name in TIMM_MODELS
 
 def list_extended_timm_models() -> List[str]:
-    return TIMM_MODELS.keys()
+    return list(TIMM_MODELS.keys())
 
 # TODO - Figure out the reason of cold start memory spike
 BATCH_SIZE_DIVISORS = {