diff --git a/src/AIMOS.ml b/src/AIMOS.ml
index 5ee489ee02556fd7948d0dc2745a16554b33fd18..2674d286c82881d248fafd85c612da36c57ee02d 100644
--- a/src/AIMOS.ml
+++ b/src/AIMOS.ml
@@ -25,19 +25,18 @@ open Base
 
 let aimos_file = Re__Core.(compile (str "%{aimos_file}"))
 
-let write_config (inputs_path : string) (models_path : string)
-  (perturbation : string) (perturbation_path : string) (out_mode : string)
-  (amplitude : Dataset.amplitude) =
+let write_config inputs_path models_path perturbation perturbation_path out_mode
+  amplitude =
   let config_file = Caml.Filename.temp_file "aimos-" ".yml" in
   let config_path = Fpath.v config_file in
   let ampli = Dataset.string_of_amplitude amplitude in
-  let transformations : Yaml.value =
+  let transformations =
     match ampli with
     | None -> `A [ `O [ ("name", `String perturbation) ] ]
     | Some a ->
       `A [ `O [ ("name", `String perturbation); ("fn_range", `String a) ] ]
   in
-  let options : Yaml.value =
+  let options =
     match perturbation_path with
     | "" ->
       `O
@@ -55,7 +54,7 @@ let write_config (inputs_path : string) (models_path : string)
           ("custom_t_path", `String ctp);
         ]
   in
-  let models : Yaml.value =
+  let models =
     match out_mode with
     | "" ->
       `A [ `O [ ("defaults", `O [ ("models_path", `String models_path) ]) ] ]
@@ -73,20 +72,15 @@ let write_config (inputs_path : string) (models_path : string)
             ];
         ]
   in
-  let full_config : Yaml.value =
-    `O [ ("options", options); ("models", models) ]
-  in
+  let full_config = `O [ ("options", options); ("models", models) ] in
   Yaml_unix.to_file_exn config_path full_config;
   config_file
 
 let build_command config_prover
   (predicate : (Language.nn_shape, string) Dataset.predicate) =
-  let dataset = predicate.dataset in
-  let inputs_path, models_path =
-    (Unix.realpath dataset, Unix.realpath predicate.model.filename)
-  in
-  let perturbation, (amplitude : Dataset.amplitude), perturbation_path, out_mode
-      =
+  let inputs_path = Unix.realpath predicate.dataset in
+  let models_path = Unix.realpath predicate.model.filename in
+  let perturbation, amplitude, perturbation_path, out_mode =
     match predicate.property with
     | Dataset.MetaRobust (_, aimos_params) ->
       ( aimos_params.perturbation,
diff --git a/src/dataset.ml b/src/dataset.ml
index 3da27245245f17bca8ba4bcddfb1994e3419064a..8b89cfe8bb9c1d042c0a60dcd52ee02e28d3d467 100644
--- a/src/dataset.ml
+++ b/src/dataset.ml
@@ -113,7 +113,6 @@ let failwith_unsupported_term t =
 let failwith_unsupported_ls ls =
   failwith (Fmt.str "Unsupported logic symbol %a" Pretty.print_ls ls)
 
-(* TODO: use amplitude type instead of 3 ints *)
 let interpret_predicate env ~on_model ~on_dataset task =
   let task = Trans.apply Introduction.introduce_premises task in
   let term = Task.task_goal_fmla task in