diff --git a/mffusion/demo/dmfal_demo.py b/mffusion/demo/dmfal_demo.py
index 5f2ddeff..09f386c6 100644
--- a/mffusion/demo/dmfal_demo.py
+++ b/mffusion/demo/dmfal_demo.py
@@ -108,7 +108,7 @@ def dmfal_test(dataset, exp_config):
     # predict
     predict_y = dmfal_model.predict(eval_inputs)
     predict_y = data_norm_manager.denormalize_output(predict_y, -1)
-    # plot_result(eval_outputs[-1], predict_y, source_shape)
+    plot_result(eval_outputs[-1], predict_y, source_shape)
 
     from mffusion.utils.performance_evaluator import performance_evaluator
     eval_result = performance_evaluator(eval_outputs[-1], predict_y, ['rmse', 'r2'])
diff --git a/mffusion/demo/fides_demo.py b/mffusion/demo/fides_demo.py
index 1e6ba665..7bacd243 100644
--- a/mffusion/demo/fides_demo.py
+++ b/mffusion/demo/fides_demo.py
@@ -155,7 +155,7 @@ def gp_model_block_test(dataset, exp_config):
         predict_y = fides_block.predict([eval_inputs[0], predict_y_mean])
         predict_y_mean = predict_y[0]
 
-    # plot_result(eval_outputs[len(train_outputs)-1], predict_y_mean, source_shape)
+    plot_result(eval_outputs[len(train_outputs)-1], predict_y_mean, source_shape)
 
     from mffusion.utils.performance_evaluator import performance_evaluator
     eval_result = performance_evaluator(eval_outputs[len(train_outputs)-1], predict_y_mean, ['rmse', 'r2'])
diff --git a/mffusion/demo/gar_demo.py b/mffusion/demo/gar_demo.py
index e6d5497f..d466c4aa 100644
--- a/mffusion/demo/gar_demo.py
+++ b/mffusion/demo/gar_demo.py
@@ -92,7 +92,7 @@ def gp_model_block_test(_num, exp_config):
     print('\n')
     gp_model_block.eval()
     predict_y = gp_model_block.predict(eval_inputs)
-    # plot_result(eval_outputs[0], predict_y, source_shape)
+    plot_result(eval_outputs[0], predict_y, source_shape)
 
     from mffusion.utils.performance_evaluator import performance_evaluator
     eval_result = performance_evaluator(eval_outputs[0], predict_y[0].mean, ['rmse', 'r2'])
diff --git a/mffusion/demo/ifc_gpt_2d_demo.py b/mffusion/demo/ifc_gpt_2d_demo.py
index c5ffa76f..7a208f65 100644
--- a/mffusion/demo/ifc_gpt_2d_demo.py
+++ b/mffusion/demo/ifc_gpt_2d_demo.py
@@ -122,14 +122,14 @@ def ifc_gpt_2d_test(dataset, exp_config):
             # mae_list_te = ifc_model.eval_mae(Xte_list, yte_list, t_list_te, t_list_tr)
             
             pred_dict = ifc_model.eval_pred(Xte_list, t_list_te, t_list_tr)
-            # plot_result(eval_outputs[-1], predict_y, source_shape)
             
 
-            from utils.performance_evaluator import performance_evaluator
+            from mffusion.utils.performance_evaluator import performance_evaluator
             src_eval_output = data_norm_manager.denormalize_output(eval_outputs[-1], -1)
             pred_output = pred_dict[ifc_dataset.fid_list_tr[-1]]
             pred_output = data_norm_manager.denormalize_output(torch.tensor(pred_output).reshape_as(eval_outputs[-1]), -1)
 
+            # plot_result(eval_outputs[-1], pred_output, source_shape)
             eval_result = performance_evaluator(src_eval_output, pred_output, ['rmse', 'r2'])
             eval_result['time'] = time.time()-start_time
             eval_result['train_sample_num'] = train_inputs[0].shape[0]
@@ -138,7 +138,7 @@ def ifc_gpt_2d_test(dataset, exp_config):
 
             scheduler.step(adjust_rmse)
 
-        #
+        # plot_result(eval_outputs[-1], pred_output, source_shape)
 
         optimizer.zero_grad()
         loss.backward()
diff --git a/mffusion/demo/ifc_ode_2d_demo.py b/mffusion/demo/ifc_ode_2d_demo.py
index fae2045b..e4b7669a 100644
--- a/mffusion/demo/ifc_ode_2d_demo.py
+++ b/mffusion/demo/ifc_ode_2d_demo.py
@@ -126,10 +126,8 @@ def ifc_gpt_2d_test(dataset, exp_config):
             # mae_list_te = ifc_model.eval_mae(Xte_list, yte_list, t_list_te, t_list_tr)
             
             pred_dict = ifc_model.eval_pred(Xte_list, t_list_te)
-            # plot_result(eval_outputs[-1], predict_y, source_shape)
             
-
-            from utils.performance_evaluator import performance_evaluator
+            from mffusion.utils.performance_evaluator import performance_evaluator
             src_eval_output = data_norm_manager.denormalize_output(eval_outputs[-1], -1)
             pred_output = pred_dict[ifc_dataset.fid_list_tr[-1]]
             pred_output = data_norm_manager.denormalize_output(torch.tensor(pred_output).reshape_as(eval_outputs[-1]), -1)
@@ -141,7 +139,7 @@ def ifc_gpt_2d_test(dataset, exp_config):
 
             scheduler.step(adjust_rmse)
 
-        #
+        # plot_result(eval_outputs[-1], pred_output, source_shape)
 
         optimizer.zero_grad()
         loss.backward()
diff --git a/mffusion/modules/gp_module/hogp.py b/mffusion/modules/gp_module/hogp.py
index 1579860c..74a51e3b 100644
--- a/mffusion/modules/gp_module/hogp.py
+++ b/mffusion/modules/gp_module/hogp.py
@@ -14,7 +14,6 @@
 from mffusion.utils import *
 from mffusion.modules.kernel import kernel_utils
 from mffusion.utils.data_utils import data_register
-from mffusion.utils.mlgp_hook import set_function_as_module_to_catch_error
 from mffusion.modules.gp_module.basic_gp_model import BASE_GP_MODEL