From b17824cba64d37c26eeccb2a6c338b77b9d60236 Mon Sep 17 00:00:00 2001
From: nabenabe0928 <shuhei.watanabe.utokyo@gmail.com>
Date: Tue, 10 Dec 2024 08:58:01 +0100
Subject: [PATCH] Fix

---
 package/samplers/hebo/sampler.py | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/package/samplers/hebo/sampler.py b/package/samplers/hebo/sampler.py
index 5c94f922..2d7d8695 100644
--- a/package/samplers/hebo/sampler.py
+++ b/package/samplers/hebo/sampler.py
@@ -99,16 +99,15 @@ def _suggest_and_transform_to_dict(
                 continue
 
             dist = search_space[name]
-            if isinstance(dist, (IntDistribution, FloatDistribution)):
-                if not dist.log and dist.step is not None:
-                    step_index = row.iloc[0]
-                    params[name] = dist.low + step_index * dist.step
-                else:
-                    params[name] = row.iloc[0]
-            elif isinstance(dist, CategoricalDistribution):
-                params[name] = row.iloc[0]
+            if (
+                isinstance(dist, (IntDistribution, FloatDistribution))
+                and not dist.log
+                and dist.step is not None
+            ):
+                step_index = row.iloc[0]
+                params[name] = dist.low + step_index * dist.step
             else:
-                assert False, f"Should not reach. Got an unknown distribution: {dist}."
+                params[name] = row.iloc[0]
 
         return params