make sure that the behavior is consistent for attrgetter

pyro-ppl · Nov 13, 2024 · 4407e60 · 4407e60
1 parent 33cd1fd
commit 4407e60
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/numpyro/util.py b/numpyro/util.py
@@ -787,7 +787,8 @@ def nested_attrgetter(*collect_fields):
     """
 
     def getter(obj):
-        return tuple(_get_nested_attr(obj, field) for field in collect_fields)
+        results = tuple(_get_nested_attr(obj, field) for field in collect_fields)
+        return results if len(collect_fields) > 1 else results[0]
 
     return getter