move copy inside get_ith_value

aleatory-science · Dec 18, 2023 · 9c18d50 · 9c18d50
1 parent a975ab3
commit 9c18d50
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/numpyro/contrib/control_flow/scan.py b/numpyro/contrib/control_flow/scan.py
@@ -26,10 +26,11 @@ def get_ith_value(site):
         ):
             return site
 
+        site = site.copy()
         site["value"] = site["value"][i]
         return site
 
-    return {k: get_ith_value(v.copy()) for k, v in replay_trace.items()}
+    return {k: get_ith_value(v) for k, v in replay_trace.items()}
 
 
 def _subs_wrapper(subs_map, i, length, site):