Update davit.py

Smartappli · Jun 25, 2024 · acde2e9 · acde2e9
1 parent f03006d
commit acde2e9
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/fl_common/models/davit.py b/fl_common/models/davit.py
@@ -21,6 +21,8 @@ def get_davit_model(davit_type, num_classes):
         "davit_large",
         "davit_huge",
         "davit_giant",
+        "davit_base_fl",
+        "davit_huge_fl",
     ]
 
     # Check if the davit_type is valid