made sure that everything is GPU compatible

ThomasHelfer · Jan 4, 2024 · d069af4 · d069af4
1 parent 4a36eaf
commit d069af4
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/GeneralRelativity/TensorAlgebra.py b/GeneralRelativity/TensorAlgebra.py
@@ -22,8 +22,8 @@ def compute_christoffel(d1_metric: torch.tensor, h_UU: torch.tensor) -> torch.te
     """
 
     chris = {
-        "LLL": torch.zeros(d1_metric.shape, dtype=d1_metric.dtype),
-        "ULL": torch.zeros(d1_metric.shape, dtype=d1_metric.dtype),
+        "LLL": torch.zeros_like(d1_metric),
+        "ULL": torch.zeros_like(d1_metric),
     }
 
     # Compute Christoffel symbols of the first kind (LLL)
@@ -61,8 +61,8 @@ def compute_christoffel_fast(
     # Initialize the output tensors
     # batch, x, y, z, i, j, dx = d1_metric.shape
     chris = {
-        "LLL": torch.zeros(d1_metric.shape, dtype=d1_metric.dtype),
-        "ULL": torch.zeros(d1_metric.shape, dtype=d1_metric.dtype),
+        "LLL": torch.zeros_like(d1_metric),
+        "ULL": torch.zeros_like(d1_metric),
     }
 
     # Compute Christoffel symbols of the first kind (LLL)

diff --git a/GeneralRelativity/Utils.py b/GeneralRelativity/Utils.py
@@ -37,6 +37,7 @@ def __getitem__(self, key: list) -> torch.tensor:
                 + (3, 3)
                 + self.tensor.shape[(self.num_index + 1) :],
                 dtype=self.tensor.dtype,
+                device=self.device,
             )
             for i in range(3):
                 for j in range(i, 3):