diff --git a/mindarmour/reliability/model_fault_injection/fault_injection.py b/mindarmour/reliability/model_fault_injection/fault_injection.py
index a6784ef..9f7d4ec 100644
--- a/mindarmour/reliability/model_fault_injection/fault_injection.py
+++ b/mindarmour/reliability/model_fault_injection/fault_injection.py
@@ -80,7 +80,7 @@ class FaultInjector:
             if param['fi_type'] not in attr:
                 msg = "'Undefined fault type', got {}.".format(self.fi_type_map[param['fi_type']])
                 LOGGER.error(TAG, msg)
-                raise AttributeError(msg)
+                raise ValueError(msg)
             if param['fi_mode'] not in ['single_layer', 'all_layer']:
                 msg = "'fault mode should be single_layer or all_layer', but got {}.".format(param['fi_mode'])
                 LOGGER.error(TAG, msg)
diff --git a/tests/ut/python/reliability/model_fault_injection/test_fault_injection.py b/tests/ut/python/reliability/model_fault_injection/test_fault_injection.py
index f99cab2..59be4d4 100644
--- a/tests/ut/python/reliability/model_fault_injection/test_fault_injection.py
+++ b/tests/ut/python/reliability/model_fault_injection/test_fault_injection.py
@@ -195,11 +195,11 @@ def test_wrong_fi_type():
     fi_size = [1]
 
     # Fault injection
-    with pytest.raises(AttributeError) as exc_info:
+    with pytest.raises(ValueError) as exc_info:
         fi = FaultInjector(model, fi_type, fi_mode, fi_size)
         _ = fi.kick_off(ds_data, ds_label, iter_times=100)
         _ = fi.metrics()
-    assert exc_info.type is AttributeError
+    assert exc_info.type is ValueError
 
 
 @pytest.mark.level0