diff --git a/hatlib/hat.py b/hatlib/hat.py
index 3c80e73..1baa78b 100644
--- a/hatlib/hat.py
+++ b/hatlib/hat.py
@@ -95,7 +95,7 @@ def generate_arg_sets_for_func(
 
     arg_sets = [generate_arg_values(parameters, dim_names_to_values) for _ in range(num_input_sets)]
 
-    return arg_sets
+    return arg_sets[0] if len(arg_sets) == 1 else arg_sets
 
 
 def generate_arg_sets_for_hat_file(hat_path):