feat(inference): mixed precision support for base encoder/coarsener

ZettaAI · Nov 8, 2023 · 7f0a270 · 7f0a270
1 parent ede7425
commit 7f0a270
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/zetta_utils/alignment/base_coarsener.py b/zetta_utils/alignment/base_coarsener.py
@@ -56,7 +56,8 @@ def __call__(self, src: torch.Tensor) -> torch.Tensor:
                     y_end = y + self.tile_size + self.tile_pad_in
                     tile = data_in[:, :, x_start:x_end, y_start:y_end]
                     if (tile != 0).sum() > 0.0:
-                        tile_result = model(tile)
+                        with torch.autocast(device_type=device):
+                            tile_result = model(tile)
                         if tile_pad_out > 0:
                             tile_result = tile_result[
                                 :, :, tile_pad_out:-tile_pad_out, tile_pad_out:-tile_pad_out

diff --git a/zetta_utils/alignment/base_encoder.py b/zetta_utils/alignment/base_encoder.py
@@ -36,7 +36,8 @@ def __call__(self, src: torch.Tensor) -> torch.Tensor:
                 raise ValueError(f"Unsupported src dtype: {src.dtype}")
 
             data_in = einops.rearrange(data_in, "C X Y Z -> Z C X Y")
-            result = model(data_in.to(device))
+            with torch.autocast(device_type=device):
+                result = model(data_in.to(device))
             result = einops.rearrange(result, "Z C X Y -> C X Y Z")
 
             # Final layer assumed to be tanh