torch backend no_sum fix

danlkv · Mar 19, 2024 · 7e0286f · 7e0286f
1 parent 6f30c94
commit 7e0286f
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/qtensor/contraction_backends/torch.py b/qtensor/contraction_backends/torch.py
@@ -140,7 +140,10 @@ def process_bucket(self, bucket, no_sum=False):
                 key=int, reverse=True
             ))
         else:
-            result_data = result_data.sum(axis=-1)
+            if not no_sum:
+                result_data = result_data.sum(axis=-1)
+            else:
+                result_data = result_data