fix calc class weight

DataCanvasIO · Feb 22, 2024 · 253e50d · 253e50d
1 parent a92c0fa
commit 253e50d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deeptables/models/deeptable.py b/deeptables/models/deeptable.py
@@ -663,7 +663,7 @@ def get_class_weight(self, y):
         #     logger.info(f'class {i}:{weight}')
 
         n = len(self.classes_)
-        class_weight = get_tool_box(y).compute_class_weight('balanced', classes=range(n), y=y)
+        class_weight = get_tool_box(y).compute_class_weight('balanced', classes=np.array(list(range(n))), y=y)
         class_weight = {k: v for k, v in zip(range(n), class_weight)}
         logger.info(f'classes weight: {class_weight}')