Replace torch._utils._accumulate with numpy.cumsum. (#3364)

pyro-ppl · May 3, 2024 · ca36025 · ca36025
1 parent 5ba6232
commit ca36025
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 6 deletions.
diff --git a/examples/contrib/mue/FactorMuE.py b/examples/contrib/mue/FactorMuE.py
@@ -86,9 +86,7 @@ def main(args):
         indices = torch.randperm(sum(data_lengths), device=device).tolist()
         dataset_train, dataset_test = [
             torch.utils.data.Subset(dataset, indices[(offset - length) : offset])
-            for offset, length in zip(
-                torch._utils._accumulate(data_lengths), data_lengths
-            )
+            for offset, length in zip(np.cumsum(data_lengths), data_lengths)
         ]
     else:
         dataset_train = dataset

diff --git a/examples/contrib/mue/ProfileHMM.py b/examples/contrib/mue/ProfileHMM.py
@@ -92,9 +92,7 @@ def main(args):
         indices = torch.randperm(sum(data_lengths), device=device).tolist()
         dataset_train, dataset_test = [
             torch.utils.data.Subset(dataset, indices[(offset - length) : offset])
-            for offset, length in zip(
-                torch._utils._accumulate(data_lengths), data_lengths
-            )
+            for offset, length in zip(np.cumsum(data_lengths), data_lengths)
         ]
     else:
         dataset_train = dataset