Merge branch 'main' of github-personal:microsoft/aurora

aurora/batch.py

@@ -135,6 +135,6 @@ def to(self, device: str | torch.device) -> "Batch":
         """Move the batch to another device."""
         return self._fmap(lambda x: x.to(device))
 
-    def float(self) -> "Batch":
-        """Convert everything to `float32`s."""
-        return self._fmap(lambda x: x.float())
+    def type(self, t: type) -> "Batch":
+        """Convert everything to type `t`."""
+        return self._fmap(lambda x: x.type(t))

aurora/model/aurora.py

@@ -111,11 +111,12 @@ def forward(self, batch: Batch) -> Batch:
         Returns:
             :class:`Batch`: Prediction for the batch.
         """
-        batch = batch.float()  # `float64`s will take up too much memory.
+        # Get the first parameter. We'll derive the data type and device from this parameter.
+        p = next(self.parameters())
+        batch = batch.type(p.dtype)
         batch = batch.normalise()
         batch = batch.crop(patch_size=self.patch_size)
-        # Assume that all parameters of the model are either on the CPU or GPU.
-        batch = batch.to(next(self.parameters()).device)
+        batch = batch.to(p.device)
 
         H, W = batch.spatial_shape
         patch_res: Int3Tuple = (

tests/test_headers.py

@@ -15,7 +15,7 @@
     relative_path = path.relative_to(_root)
 
     # Ignore a possible virtual environment.
-    if str(relative_path.parents[-2]) in {"venv"}:
+    if len(relative_path.parents) >= 2 and str(relative_path.parents[-2]) in {"venv"}:
         continue
 
     # Ignore the automatically generated version file.