From e583d85ff932e4f3a9c59c141a3a812746772147 Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Wed, 6 May 2026 12:51:51 -0700 Subject: [PATCH] fix dtypes for torch --- ScaFFold/utils/data_types.py | 5 ++++- ScaFFold/utils/trainer.py | 10 ++++++---- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/ScaFFold/utils/data_types.py b/ScaFFold/utils/data_types.py index b555811..ef1515d 100644 --- a/ScaFFold/utils/data_types.py +++ b/ScaFFold/utils/data_types.py @@ -19,7 +19,10 @@ # Masks are values 0 <= x <= n_categories MASK_DTYPE = np.uint16 # Volumes/img are 0 <= x <= 1 -VOLUME_DTYPE = np.float32 +VOLUME_DTYPE_NAME = "float32" +VOLUME_NP_DTYPE = getattr(np, VOLUME_DTYPE_NAME) +VOLUME_TORCH_DTYPE = getattr(torch, VOLUME_DTYPE_NAME) +VOLUME_DTYPE = VOLUME_NP_DTYPE # Shared AMP dtype selection for torch.autocast. AMP_DTYPE = torch.bfloat16 diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 1a1d2e0..3add912 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -30,7 +30,7 @@ from ScaFFold.utils.checkpointing import CheckpointManager from ScaFFold.utils.data_loading import FractalDataset, SpatialShardSpec -from ScaFFold.utils.data_types import AMP_DTYPE, VOLUME_DTYPE +from ScaFFold.utils.data_types import AMP_DTYPE, VOLUME_TORCH_DTYPE from ScaFFold.utils.dice_score import compute_sharded_dice from ScaFFold.utils.distributed import get_local_rank, get_world_rank, get_world_size @@ -436,7 +436,7 @@ def warmup(self): images = images.to( device=self.device, - dtype=VOLUME_DTYPE, + dtype=VOLUME_TORCH_DTYPE, memory_format=torch.channels_last_3d, non_blocking=True, ) @@ -611,7 +611,7 @@ def train(self): begin_code_region("image_to_device") images = images.to( device=self.device, - dtype=VOLUME_DTYPE, + dtype=VOLUME_TORCH_DTYPE, memory_format=torch.channels_last_3d, # NDHWC (channels last) vs NCDHW (channels first) non_blocking=True, ) @@ -749,7 +749,9 @@ def train(self): self.config.n_categories, self.config._parallel_strategy, ) - dice_info = torch.tensor([dice_sum, numsamples], dtype=VOLUME_DTYPE) + dice_info = torch.tensor( + [dice_sum, numsamples], dtype=VOLUME_TORCH_DTYPE + ) if self.config.dist: dice_info = dice_info.to(device=self.device) torch.distributed.all_reduce(