add moving results to host memory after contraction using torch

danlkv · Apr 12, 2024 · fdbff76 · fdbff76
1 parent 447a13d
commit fdbff76
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/qtensor/contraction_backends/torch.py b/qtensor/contraction_backends/torch.py
@@ -551,4 +551,6 @@ def get_result_data(self, result):
             d = result.data.reshape(self._get_index_sizes(*result.indices))
         else:
             d = result.data
+        # move to cpu
+        d = d.cpu()
         return torch.permute(d, tuple(reversed(range(d.ndim))))