diff --git a/src/quantem/core/ml/inr.py b/src/quantem/core/ml/inr.py
index 4c197344..f4801705 100644
--- a/src/quantem/core/ml/inr.py
+++ b/src/quantem/core/ml/inr.py
@@ -111,18 +111,20 @@ def _build(self) -> None:
         self.net = nn.Sequential(*net_list)
 
         if self.winner_initialization:
-            if type(self.winner_initialization) is int:
-                rng = torch.Generator()
-                rng.manual_seed(self.winner_initialization)
-            else:
-                rng = torch.Generator()
-                rng.manual_seed(42)
+            seed = self.winner_initialization if type(self.winner_initialization) is int else 42
+            rng = torch.Generator()
+            rng.manual_seed(seed)
+            # torch.randn_like ignores generators, so the noise must come from
+            # torch.randn with the seeded generator -- otherwise the "winner" seed
+            # silently has no effect and the perturbation is not reproducible.
             with torch.no_grad():
-                self.net[0].linear.weight += (  # type: ignore[reportAttributeAccessIssue]
-                    torch.randn_like(self.net[0].linear.weight) * 5 / self.first_omega_0  # type:ignore
-                )
-                self.net[1].linear.weight += (  # type: ignore[reportAttributeAccessIssue]
-                    torch.randn_like(self.net[1].linear.weight) * 0.1 / self.hidden_omega_0  # type:ignore
+                w0 = self.net[0].linear.weight  # type: ignore[reportAttributeAccessIssue]
+                w0 += torch.randn(w0.shape, generator=rng, dtype=w0.dtype) * 5 / self.first_omega_0
+                w1 = self.net[1].linear.weight  # type: ignore[reportAttributeAccessIssue]
+                w1 += (
+                    torch.randn(w1.shape, generator=rng, dtype=w1.dtype)
+                    * 0.1
+                    / self.hidden_omega_0
                 )
 
     def forward(self, coords: torch.Tensor) -> torch.Tensor:
diff --git a/src/quantem/core/ml/models/kplanes.py b/src/quantem/core/ml/models/kplanes.py
index cdf55261..93dad498 100644
--- a/src/quantem/core/ml/models/kplanes.py
+++ b/src/quantem/core/ml/models/kplanes.py
@@ -169,6 +169,7 @@ class KPlanes(PPLR, TensorDecompositionModel):
     """
     K-Planes model adapted from Fridovich-Keil et al., https://arxiv.org/abs/2301.10241
     """
+
     def __init__(
         self,
         # Grid parameters
@@ -199,6 +200,15 @@ def __init__(
         self.concat_features = concat_features
         self.density_activation = density_activation
 
+        # All three planes share one (3, C, res[1], res[0]) tensor, which ignores
+        # res[2]: an anisotropic resolution would silently give the XZ/YZ planes the
+        # wrong grid along z. Refuse it rather than misallocate.
+        if len(set(self.resolution)) != 1:
+            raise ValueError(
+                f"KPlanes currently requires an isotropic resolution, got {self.resolution}; "
+                "the plane grids are allocated as (res[1], res[0]) for all three axis pairs."
+            )
+
         self.grids = nn.ParameterList()
         self.feature_dim = 0
         for res_mult in self.multiscale_res_multipliers:
@@ -320,14 +330,18 @@ def interpolate_ms_features_tilted(
     rotated = torch.einsum("tij,bj->tbi", rotation_matrices, pts)
 
     # Build (T, 3, B, 2) coords for planes XY, ZX, YZ in one shot.
-    # index_select is faster and cleaner than advanced indexing with python lists.
-    # Plane axis layout: XY=(0,1), ZX=(2,0), YZ=(1,2)
-    idx = torch.tensor([[0, 1], [2, 0], [1, 2]], device=pts.device)  # (3, 2)
-    # rotated: (T, B, 3) -> gather along last dim with idx (3, 2)
-    # Result: (T, 3, B, 2)
-    coords = (
-        rotated.unsqueeze(1).expand(T, 3, B, 3).gather(-1, idx.view(1, 3, 1, 2).expand(T, 3, B, 2))
-    )
+    # Plane axis layout: XY=(0,1), ZX=(2,0), YZ=(1,2). Stacking views avoids the
+    # per-call index-tensor allocation (a host-to-device copy) and the
+    # (T, 3, B, 3) expand+gather this used to do.
+    x, y, z = rotated.unbind(-1)  # each (T, B)
+    coords = torch.stack(
+        (
+            torch.stack((x, y), dim=-1),
+            torch.stack((z, x), dim=-1),
+            torch.stack((y, z), dim=-1),
+        ),
+        dim=1,
+    )  # (T, 3, B, 2)
 
     # Flatten (T, 3) -> 3*T so it matches grid's first dim, and add the H_out=1 axis
     coord_tensor = coords.reshape(3 * T, B, 1, 2)  # (3T, B, 1, 2)
diff --git a/src/quantem/core/ml/optimizer_mixin.py b/src/quantem/core/ml/optimizer_mixin.py
index ece31dff..af125391 100644
--- a/src/quantem/core/ml/optimizer_mixin.py
+++ b/src/quantem/core/ml/optimizer_mixin.py
@@ -294,14 +294,15 @@ class Plateau:
         _name: str = "plateau"
 
         def params(self, base_LR: float, num_iter: int | None = None) -> dict:
-            if self.min_lr is None:
-                self.min_lr = self.min_lr_factor * base_LR
+            # Derived values stay local: params() must not mutate the dataclass, or a
+            # shared instance bakes in the first optimizer's base LR for every later one.
+            min_lr = self.min_lr if self.min_lr is not None else self.min_lr_factor * base_LR
             return {
                 "mode": self.mode,
                 "factor": self.factor,
                 "patience": self.patience,
                 "threshold": self.threshold,
-                "min_lr": self.min_lr,
+                "min_lr": min_lr,
                 "cooldown": self.cooldown,
             }
 
@@ -332,13 +333,12 @@ def params(self, base_LR: float, num_iter: int | None = None) -> dict:
             if effective_num_iter is None:
                 raise ValueError("num_iter must be set if num_iter is not provided")
 
-            self.num_iter = effective_num_iter
-
+            gamma = self.gamma
             if self.factor is not None:
-                self.gamma = self.factor ** (1.0 / effective_num_iter)
+                gamma = self.factor ** (1.0 / effective_num_iter)
 
             return {
-                "gamma": self.gamma,
+                "gamma": gamma,
             }
 
     @dataclass
@@ -385,13 +385,11 @@ class Cyclic:
         _name: str = "cyclic"
 
         def params(self, base_LR: float, num_iter: int | None = None) -> dict:
-            if self.base_lr is None:
-                self.base_lr = self.base_lr_factor * base_LR
-            if self.max_lr is None:
-                self.max_lr = self.max_lr_factor * base_LR
+            base_lr = self.base_lr if self.base_lr is not None else self.base_lr_factor * base_LR
+            max_lr = self.max_lr if self.max_lr is not None else self.max_lr_factor * base_LR
             return {
-                "base_lr": self.base_lr,
-                "max_lr": self.max_lr,
+                "base_lr": base_lr,
+                "max_lr": max_lr,
                 "step_size_up": self.step_size_up,
                 "step_size_down": self.step_size_down,
                 "mode": self.mode,
@@ -426,12 +424,11 @@ def params(self, base_LR: float, num_iter: int | None = None) -> dict:
                 raise ValueError(
                     "total_iters must be set if num_iter is not provided"
                 )  # Should never be reached
-            if self.total_iters is None:
-                self.total_iters = num_iter
+            total_iters = self.total_iters if self.total_iters is not None else num_iter
             return {
                 "start_factor": self.start_factor,
                 "end_factor": self.end_factor,
-                "total_iters": self.total_iters,
+                "total_iters": total_iters,
             }
 
     @dataclass
@@ -459,10 +456,9 @@ def params(self, base_LR: float, num_iter: int | None = None) -> dict:
                 raise ValueError(
                     "T_max must be set if num_iter is not provided"
                 )  # Should never be reached
-            if self.T_max is None:
-                self.T_max = num_iter
+            T_max = self.T_max if self.T_max is not None else num_iter
             return {
-                "T_max": self.T_max,
+                "T_max": T_max,
                 "eta_min": self.eta_min,
             }
 
@@ -572,7 +568,9 @@ def _normalize_optimizer_params(
         if isinstance(params, OptimizerParamsType):
             return {self.DEFAULT_OPTIMIZER_KEY: params}
         if not isinstance(params, dict):
-            raise TypeError(f"optimizer_params must be OptimizerParamsType or dict, got {type(params)}")
+            raise TypeError(
+                f"optimizer_params must be OptimizerParamsType or dict, got {type(params)}"
+            )
         # Single optimizer as dict shorthand, e.g. {"name": "adam", "lr": 1e-3}
         if self._is_single_optimizer_dict(params):
             return {self.DEFAULT_OPTIMIZER_KEY: OptimizerParams.parse_dict(d=params)}
@@ -597,7 +595,9 @@ def scheduler_params(self, params: SchedulerParamsType | dict):
         if isinstance(params, dict):
             params = SchedulerParams.parse_dict(d=params)
         if not isinstance(params, SchedulerParamsType):
-            raise TypeError(f"scheduler parameters must be a SchedulerParamsType, got {type(params)}")
+            raise TypeError(
+                f"scheduler parameters must be a SchedulerParamsType, got {type(params)}"
+            )
         self._scheduler_params = params
 
     @abstractmethod
@@ -688,7 +688,9 @@ def _build_optimizer(self, opt_params, param_groups) -> "torch.optim.Optimizer":
                 raise NotImplementedError(f"Unknown optimizer type: {opt_params}")
 
     def set_scheduler(
-        self, scheduler_params: SchedulerParamsType | dict | None = None, num_iter: int | None = None
+        self,
+        scheduler_params: SchedulerParamsType | dict | None = None,
+        num_iter: int | None = None,
     ) -> None:
         """Set the scheduler for this model."""
         if scheduler_params is not None:
diff --git a/src/quantem/tomography/dataset_models.py b/src/quantem/tomography/dataset_models.py
index 2255636c..8b6cb8ba 100644
--- a/src/quantem/tomography/dataset_models.py
+++ b/src/quantem/tomography/dataset_models.py
@@ -189,6 +189,12 @@ def __init__(
         if type(tilt_angles) is not torch.Tensor:
             tilt_angles = torch.from_numpy(tilt_angles)
         max_val = torch.quantile(tilt_stack, 0.95)
+        # A sparse stack (>95% zeros) has a zero 95th quantile; dividing by it
+        # would turn the targets into inf/NaN and poison the first backward.
+        if max_val <= 0:
+            max_val = tilt_stack.abs().max()
+        if max_val <= 0:
+            raise ValueError("tilt_stack is all zeros; cannot normalize.")
 
         # Tilt stack normalization
         tilt_stack = tilt_stack / max_val
@@ -240,6 +246,21 @@ def get_optimization_parameters(self) -> dict[str, list[torch.Tensor]]:
         """
         return {self.DEFAULT_OPTIMIZER_KEY: list(self.parameters())}
 
+    def _materialize_pose_parameters(self, device: str | torch.device):
+        """Create the learnable pose parameters, or move the existing ones.
+
+        Once the parameters exist, their *current* (possibly trained) values are
+        moved; the initial-value buffers are only used on first materialization.
+        Rebuilding from the buffers on every call silently reset learned poses
+        whenever the dataset changed device (e.g. ``from_file(...).to(device)``).
+        """
+        z1 = self._z1_params.data if hasattr(self, "_z1_params") else self._z1_angles
+        z3 = self._z3_params.data if hasattr(self, "_z3_params") else self._z3_angles
+        shifts = self._shifts_params.data if hasattr(self, "_shifts_params") else self._shifts
+        self._z1_params = nn.Parameter(z1.detach().to(device))
+        self._z3_params = nn.Parameter(z3.detach().to(device))
+        self._shifts_params = nn.Parameter(shifts.detach().to(device))
+
     # --- Forward pass ---
     @abstractmethod
     def forward(
@@ -302,12 +323,11 @@ def reference_tilt_idx(self, reference_tilt_idx: int):
 
     @property
     def learnable_tilts(self) -> int:
+        # Derived from the tilt series (all tilts minus the fixed reference); there is
+        # deliberately no setter -- the old one wrote a private attribute this getter
+        # never read, so assignments appeared to succeed while doing nothing.
         return self.tilt_angles.shape[0] - 1
 
-    @learnable_tilts.setter
-    def learnable_tilts(self, learnable_tilts: int):
-        self._learnable_tilts = learnable_tilts
-
     @property
     def z1_params(self) -> torch.nn.Parameter:
         return self._z1_params
@@ -429,9 +449,7 @@ def to(self, device: str | torch.device):
         self.tilt_stack = self.tilt_stack.to(device)
         self.tilt_angles = self.tilt_angles.to(device)
 
-        self._z1_params = nn.Parameter(self._z1_angles.to(device))
-        self._z3_params = nn.Parameter(self._z3_angles.to(device))
-        self._shifts_params = nn.Parameter(self._shifts.to(device))
+        self._materialize_pose_parameters(device)
 
         self._z1_ref = self._z1_ref.to(device)
         self._z3_ref = self._z3_ref.to(device)
@@ -440,6 +458,96 @@ def to(self, device: str | torch.device):
         self.device = device
 
 
+class DeviceBatchSampler:
+    """Epoch iterator that builds INR training batches directly on a device.
+
+    Replaces the per-pixel DataLoader path: the tilt stack and angles are
+    made resident on ``device`` once, so producing a batch is index
+    arithmetic plus two tensor lookups instead of ``batch_size`` Python
+    ``__getitem__`` calls, a collate, and a host-to-device copy per step.
+    On a GPU this removes the CPU dataloader bottleneck entirely.
+
+    Yields the same batch dicts as ``TomographyINRDataset.__getitem__``
+    under a DataLoader collate (``projection_idx``, ``pixel_i``,
+    ``pixel_j``, ``phi``, ``target_value``), with the train loader's
+    ``drop_last=True`` semantics.
+
+    Distributed runs: pass ``rank``/``world_size`` and every rank derives
+    the *same* epoch permutation from ``seed + epoch`` (CPU generator, so
+    it is identical across ranks and reproducible), then takes an
+    equal-size contiguous shard — equal so per-rank batch counts match and
+    DDP gradient sync cannot hang on a ragged tail. The training loop's
+    ``sampler.set_epoch(epoch)`` drives reshuffling, exactly like
+    ``DistributedSampler``; without ``set_epoch`` the epoch advances
+    automatically on each ``__iter__``.
+    """
+
+    def __init__(
+        self,
+        dset: "TomographyINRDataset",
+        batch_size: int,
+        device: torch.device | str,
+        indices: torch.Tensor | None = None,
+        shuffle: bool = True,
+        rank: int = 0,
+        world_size: int = 1,
+        seed: int = 0,
+    ):
+        self.batch_size = batch_size
+        self.device = torch.device(device)
+        self.shuffle = shuffle
+        self.rank = rank
+        self.world_size = world_size
+        self.seed = seed
+        self._epoch = 0
+        self._stack = dset.tilt_stack.to(self.device)
+        self._angles = dset.tilt_angles.to(self.device)
+        # __getitem__ decodes flat indices with shape[1] for both rows and
+        # columns; replicate it exactly.
+        self._s1 = dset.tilt_stack.shape[1]
+        self._s2 = dset.tilt_stack.shape[2]
+        if indices is None:
+            indices = torch.arange(len(dset), dtype=torch.int64)
+        self._indices = indices.to(self.device)
+        self._per_rank = len(self._indices) // world_size
+
+    def set_epoch(self, epoch: int) -> None:
+        """Set the epoch used to seed this epoch's shared permutation."""
+        self._epoch = epoch
+
+    def __len__(self) -> int:
+        return self._per_rank // self.batch_size  # drop_last=True
+
+    def _epoch_shard(self) -> torch.Tensor:
+        idx = self._indices
+        if self.shuffle:
+            g = torch.Generator()
+            g.manual_seed(self.seed + self._epoch)
+            perm = torch.randperm(len(idx), generator=g).to(self.device)
+            idx = idx[perm]
+            self._epoch += 1  # auto-advance; set_epoch overrides per epoch
+        if self.world_size > 1:
+            idx = idx[self.rank * self._per_rank : (self.rank + 1) * self._per_rank]
+        return idx
+
+    def __iter__(self):
+        idx = self._epoch_shard()
+        per_proj = self._s1 * self._s2
+        for k in range(len(self)):
+            sel = idx[k * self.batch_size : (k + 1) * self.batch_size]
+            proj = sel // per_proj
+            rem = sel - proj * per_proj
+            pixel_i = rem // self._s1
+            pixel_j = rem - pixel_i * self._s1
+            yield {
+                "projection_idx": proj,
+                "pixel_i": pixel_i,
+                "pixel_j": pixel_j,
+                "phi": self._angles[proj],
+                "target_value": self._stack[proj, pixel_i, pixel_j],
+            }
+
+
 class TomographyINRDataset(TomographyDatasetConstraints, Dataset):
     """
     Dataset class for INR-based tomography.
@@ -550,37 +658,37 @@ def transform_batch_rays(
         shift_x_norm = (shifts[:, 0:1] * sampling_rate * 2) / (N - 1)
         shift_y_norm = (shifts[:, 1:2] * sampling_rate * 2) / (N - 1)
 
-        rays_x = rays[:, :, 0] - shift_x_norm
-        rays_y = rays[:, :, 1] - shift_y_norm
-        rays_z = rays[:, :, 2]
-
-        theta = torch.deg2rad(-z3).view(-1, 1)
-        cos_t = torch.cos(theta)
-        sin_t = torch.sin(theta)
-
-        rays_x_rot1 = cos_t * rays_x - sin_t * rays_y
-        rays_y_rot1 = sin_t * rays_x + cos_t * rays_y
-        rays_z_rot1 = rays_z
-
-        theta = torch.deg2rad(x).view(-1, 1)
-        cos_t = torch.cos(theta)
-        sin_t = torch.sin(theta)
-
-        rays_x_rot2 = rays_x_rot1
-        rays_y_rot2 = cos_t * rays_y_rot1 - sin_t * rays_z_rot1
-        rays_z_rot2 = sin_t * rays_y_rot1 + cos_t * rays_z_rot1
-
-        theta = torch.deg2rad(-z1).view(-1, 1)
-        cos_t = torch.cos(theta)
-        sin_t = torch.sin(theta)
-
-        rays_x_final = cos_t * rays_x_rot2 - sin_t * rays_y_rot2
-        rays_y_final = sin_t * rays_x_rot2 + cos_t * rays_y_rot2
-        rays_z_final = rays_z_rot2
-
-        transformed_rays = torch.stack([rays_x_final, rays_y_final, rays_z_final], dim=2)
+        shifted = torch.stack(
+            [rays[:, :, 0] - shift_x_norm, rays[:, :, 1] - shift_y_norm, rays[:, :, 2]],
+            dim=2,
+        )
 
-        return transformed_rays
+        # Compose the three Euler rotations Rz(-z1) @ Rx(x) @ Rz(-z3) into a single
+        # (B, 3, 3) matrix and apply it with one batched matmul, instead of nine
+        # elementwise passes over the full (B, S) ray tensors.
+        a = torch.deg2rad(-z3).view(-1)
+        b = torch.deg2rad(x).view(-1)
+        g = torch.deg2rad(-z1).view(-1)
+        zero = torch.zeros_like(a)
+        one = torch.ones_like(a)
+
+        cos_a, sin_a = torch.cos(a), torch.sin(a)
+        cos_b, sin_b = torch.cos(b), torch.sin(b)
+        cos_g, sin_g = torch.cos(g), torch.sin(g)
+
+        rot_a = torch.stack(
+            [cos_a, -sin_a, zero, sin_a, cos_a, zero, zero, zero, one], dim=-1
+        ).view(-1, 3, 3)
+        rot_b = torch.stack(
+            [one, zero, zero, zero, cos_b, -sin_b, zero, sin_b, cos_b], dim=-1
+        ).view(-1, 3, 3)
+        rot_g = torch.stack(
+            [cos_g, -sin_g, zero, sin_g, cos_g, zero, zero, zero, one], dim=-1
+        ).view(-1, 3, 3)
+
+        rot = rot_g @ rot_b @ rot_a  # (B, 3, 3)
+
+        return shifted @ rot.transpose(1, 2)
 
     @staticmethod
     @torch.compile(mode="reduce-overhead")
@@ -611,13 +719,16 @@ def __getitem__(
         projection_idx = actual_idx // (self.tilt_stack.shape[1] * self.tilt_stack.shape[2])
         remaining = actual_idx % (self.tilt_stack.shape[1] * self.tilt_stack.shape[2])
 
-        pixel_i = remaining // self.tilt_stack.shape[1]
-        pixel_j = remaining % self.tilt_stack.shape[1]
+        pixel_i = remaining // self.tilt_stack.shape[2]
+        pixel_j = remaining % self.tilt_stack.shape[2]
 
+        # Plain ints for the index fields: default_collate builds one int64 tensor per
+        # batch either way, but wrapping each index in torch.tensor() here allocates
+        # three scalar tensors per item on the dataloader hot path.
         return {
-            "projection_idx": torch.tensor(projection_idx),
-            "pixel_i": torch.tensor(pixel_i),
-            "pixel_j": torch.tensor(pixel_j),
+            "projection_idx": projection_idx,
+            "pixel_i": pixel_i,
+            "pixel_j": pixel_j,
             "phi": self.tilt_angles[projection_idx],  # tensor
             "target_value": self.tilt_stack[projection_idx, pixel_i, pixel_j],  # tensor
         }
@@ -628,13 +739,10 @@ def __len__(
         """
         Returns the number of pixels in the tilt stack.
         """
-        N = max(self.tilt_stack.shape)
-        return self.tilt_stack.shape[0] * N * N
+        return self.tilt_stack.shape[0] * self.tilt_stack.shape[1] * self.tilt_stack.shape[2]
 
     def to(self, device: torch.device | str):
-        self._z1_params = nn.Parameter(self._z1_angles.to(device))
-        self._z3_params = nn.Parameter(self._z3_angles.to(device))
-        self._shifts_params = nn.Parameter(self._shifts.to(device))
+        self._materialize_pose_parameters(device)
 
         self._z1_ref = self._z1_ref.to(device)
         self._z3_ref = self._z3_ref.to(device)
@@ -691,6 +799,13 @@ def __init__(
         else:
             data_quantile = torch.quantile(data, 0.95)
 
+        # Same guard as TomographyDatasetBase: a >95%-zero target has a zero
+        # 95th quantile and would normalize to inf/NaN.
+        if data_quantile <= 0:
+            data_quantile = data.abs().max()
+        if data_quantile <= 0:
+            raise ValueError("pretrain_target is all zeros; cannot normalize.")
+
         data = data / data_quantile
         data = torch.permute(data, (0, 3, 2, 1))
         # data = torch.flip(data, dims=(2,))
diff --git a/src/quantem/tomography/object_models.py b/src/quantem/tomography/object_models.py
index 817e7ed5..56e94740 100644
--- a/src/quantem/tomography/object_models.py
+++ b/src/quantem/tomography/object_models.py
@@ -421,12 +421,10 @@ def apply_hard_constraints(
 
     def apply_soft_constraints(self, ctx: ReconstructionContext) -> torch.Tensor:
         assert ctx.obj is not None, "ObjectPixelated requires ctx.obj to be set"
-        soft_loss = torch.tensor(
-            0.0, device=ctx.obj.device, dtype=ctx.obj.dtype, requires_grad=True
-        )
+        soft_loss = torch.tensor(0.0, device=ctx.obj.device, dtype=ctx.obj.dtype)
         if self.constraints.tv_vol > 0:
             tv_loss = self.get_tv_loss(ctx)
-            soft_loss += tv_loss
+            soft_loss = soft_loss + tv_loss
         return soft_loss
 
     # --- Forward method ---
@@ -542,7 +540,8 @@ def apply_soft_constraints(
         self,
         ctx: ReconstructionContext,
     ) -> torch.Tensor:
-        soft_loss = torch.tensor(0.0, device=ctx.coords.device)
+        device = ctx.coords.device if ctx.coords is not None else self._device
+        soft_loss = torch.tensor(0.0, device=device)
         if self.constraints.tv_vol > 0:
             assert ctx.coords is not None, (
                 "coords must be provided for INR object model to compute the TV loss"
@@ -662,7 +661,12 @@ def forward(self, coords: Optional[torch.Tensor] = None) -> torch.Tensor:
         if all_densities.dim() > 1:
             all_densities = all_densities.squeeze(-1)
         valid_mask = (
-            (coords[:, 0] >= -1) & (coords[:, 0] <= 1) & (coords[:, 1] >= -1) & (coords[:, 1] <= 1)
+            (coords[:, 0] >= -1)
+            & (coords[:, 0] <= 1)
+            & (coords[:, 1] >= -1)
+            & (coords[:, 1] <= 1)
+            & (coords[:, 2] >= -1)
+            & (coords[:, 2] <= 1)
         ).float()
 
         if all_densities.dim() > 1:
@@ -774,7 +778,7 @@ def create_volume(self, return_vol: bool = False):
             coords_1d = torch.linspace(-1, 1, N)
             x, y, z = torch.meshgrid(coords_1d, coords_1d, coords_1d, indexing="ij")
             inputs = torch.stack([x, y, z], dim=-1).reshape(-1, 3)
-            model = self.model.module if isinstance(self.model, nn.DataParallel) else self.model
+            model = _unwrap(self.model)
 
             inference_batch_size = 5 * N * N
             total_samples = N**3
@@ -917,7 +921,7 @@ def apply_soft_constraints(self, ctx: ReconstructionContext) -> torch.Tensor:
         soft_loss = torch.tensor(
             0.0, device=ctx.pred.device if ctx.pred is not None else self.device
         )
-        if self.constraints.tv_vol > 0:
+        if self.constraints.tv_vol > 0 or self.constraints.tv_plane > 0:
             assert ctx.coords is not None, "Coordinates must be provided for TV loss"
             assert ctx.pred is not None, "Prediction must be provided for TV loss"
             soft_loss += self.get_tv_loss(ctx)
@@ -943,18 +947,20 @@ def get_tv_loss(self, ctx: ReconstructionContext) -> torch.Tensor:
         assert ctx.coords is not None, "Coordinates must be provided for TV loss"
         assert ctx.pred is not None, "Prediction must be provided for TV loss"
         tv_loss = torch.tensor(0.0, device=ctx.pred.device)
-        tv_loss += self._get_plane_tv_loss()
-        tv_loss += self.get_volume_tv_loss(ctx.coords)
+        if self.constraints.tv_plane > 0:
+            tv_loss = tv_loss + self._get_plane_tv_loss()
+        if self.constraints.tv_vol > 0:
+            tv_loss = tv_loss + self.get_volume_tv_loss(ctx.coords)
         return tv_loss
 
     def _get_plane_tv_loss(self) -> torch.Tensor:
         """
         Gets the total-variation across the planes.
         """
-        is_tilted = self.model.tilted
+        model = _unwrap(self.model)
+        is_tilted = model.tilted
         per_level = []
 
-        model = _unwrap(self.model)
         for p in model.grids:
             # p: (3*T, C, H, W) for TILTED, (3, C, H, W) for KPlanes
             dh = (p[:, :, 1:, :] - p[:, :, :-1, :]).pow(2).mean(dim=(1, 2, 3))
@@ -962,7 +968,7 @@ def _get_plane_tv_loss(self) -> torch.Tensor:
             per_plane = dh + dw  # (3*T,) or (3,)
 
             if is_tilted:
-                T = self.model.T
+                T = model.T
                 per_rotation = per_plane.view(T, 3).sum(dim=1)  # sum 3 planes per rotation
                 level_tv = per_rotation.mean()  # avg across rotations
             else:
@@ -985,25 +991,25 @@ def get_volume_tv_loss(self, coords: torch.Tensor) -> torch.Tensor:
         model = _unwrap(self.model)
         h = 2.0 / min(model.resolution)
 
-        pred = model(tv_coords)
-        if isinstance(pred, tuple):
-            pred = pred[0]
-        if pred.dim() == 1:
-            pred = pred.unsqueeze(-1)  # (N, 1)
-
-        grads = []
-        for axis in range(3):
-            offset = torch.zeros(3, device=tv_coords.device)
-            offset[axis] = h
-            shifted_pred = self.model(tv_coords + offset)
-            if isinstance(shifted_pred, tuple):
-                shifted_pred = shifted_pred[0]
-            if shifted_pred.dim() == 1:
-                shifted_pred = shifted_pred.unsqueeze(-1)
-            grads.append((shifted_pred - pred) / h)  # (N, 1)
-
-        grad_stack = torch.stack(grads, dim=-1)  # (N, C, 3)
-        grad_norm = torch.norm(grad_stack, dim=-1)  # (N, C)
+        # Evaluate the base points and the three axis-shifted copies in a single
+        # batched forward (4N points) instead of 4 sequential model calls.
+        offsets = h * torch.eye(3, device=tv_coords.device, dtype=tv_coords.dtype)  # (3, 3)
+        all_coords = torch.cat(
+            [tv_coords, (tv_coords.unsqueeze(0) + offsets.unsqueeze(1)).reshape(-1, 3)]
+        )  # (4N, 3)
+
+        all_pred = model(all_coords)
+        if isinstance(all_pred, tuple):
+            all_pred = all_pred[0]
+        if all_pred.dim() == 1:
+            all_pred = all_pred.unsqueeze(-1)  # (4N, 1)
+
+        n = tv_coords.shape[0]
+        pred = all_pred[:n]  # (N, C)
+        shifted_pred = all_pred[n:].view(3, n, -1)  # (3, N, C)
+
+        grad_stack = (shifted_pred - pred.unsqueeze(0)) / h  # (3, N, C)
+        grad_norm = torch.norm(grad_stack, dim=0)  # (N, C)
 
         return self.constraints.tv_vol * grad_norm.mean()
 
diff --git a/src/quantem/tomography/tomography.py b/src/quantem/tomography/tomography.py
index f86095ad..ad2e65ce 100644
--- a/src/quantem/tomography/tomography.py
+++ b/src/quantem/tomography/tomography.py
@@ -17,6 +17,7 @@
     DatasetConstraintParams,
     DatasetConstraintsType,
     DatasetModelType,
+    DeviceBatchSampler,
     TomographyINRDataset,
     TomographyPixDataset,
 )
@@ -145,14 +146,7 @@ def reconstruct(
                     self.scheduler_params = scheduler_params
                     self.set_schedulers(self.scheduler_params, num_iter=num_iter)
 
-            self.dataloader, self.sampler, self.val_dataloader, self.val_sampler = (
-                self.setup_dataloader(
-                    self.dset,
-                    batch_size,
-                    num_workers=num_workers,
-                    val_fraction=val_fraction,
-                )
-            )
+            self._setup_recon_dataloaders(batch_size, num_workers, val_fraction)
 
         # Type check for INR-based reconstruction
         if not isinstance(self.dset, TomographyINRDataset):
@@ -285,10 +279,14 @@ def reconstruct(
                     val_loss = torch.tensor(0.0, device=self.device)
 
                     for batch in self.val_dataloader:
+                        # Match the training pass (enabled=False): bf16 autocast
+                        # breaks the so3 pose solve (lu_factor has no BFloat16
+                        # kernel) and would make the val loss inconsistent with
+                        # the fp32 training loss it is compared to.
                         with torch.autocast(
                             device_type=self.device.type,
                             dtype=torch.bfloat16,
-                            enabled=True,
+                            enabled=False,
                         ):
                             all_coords = self.dset.get_coords(batch, N, curr_num_samples_per_ray)
 
@@ -411,6 +409,49 @@ def _rebuild_dataloader(self, batch_size: int, num_workers: int, val_fraction: f
         """
         Rebuilds the dataloader due to persistent workers error when reloading the object.
         """
+        self._setup_recon_dataloaders(batch_size, num_workers, val_fraction)
+
+    def _setup_recon_dataloaders(self, batch_size: int, num_workers: int, val_fraction: float):
+        """Build the train/val batch iterators.
+
+        INR datasets use ``DeviceBatchSampler`` — batches are built with
+        tensor ops on the compute device from a device-resident tilt stack,
+        removing the per-pixel ``__getitem__`` / collate / H2D-copy
+        dataloader bottleneck (``num_workers`` is ignored on this path).
+        Distributed runs shard the same seeded epoch permutation across
+        ranks (DistributedSampler semantics; the loop's ``set_epoch`` drives
+        reshuffling). Non-INR datasets keep the DataLoader path.
+        """
+        if isinstance(self.dset, TomographyINRDataset):
+            n = len(self.dset)
+            n_val = int(n * val_fraction)
+            # Fixed-seed split: identical across DDP ranks (no train/val
+            # leakage between ranks) and stable across save/reload, so a
+            # resumed run keeps validating on the same held-out pixels.
+            split_gen = torch.Generator()
+            split_gen.manual_seed(0)
+            perm = torch.randperm(n, generator=split_gen)
+            ddp = dict(rank=self.global_rank, world_size=self.world_size)
+            self.dataloader = DeviceBatchSampler(
+                self.dset, batch_size, self.device, indices=perm[n_val:], **ddp
+            )
+            # The val sampler keeps its own device-resident copy of the tilt
+            # stack; acceptable, since val_fraction > 0 is the rare case.
+            val = (
+                DeviceBatchSampler(
+                    self.dset, batch_size, self.device, indices=perm[:n_val], shuffle=False, **ddp
+                )
+                if n_val > 0
+                else None
+            )
+            # A per-rank val shard smaller than one batch would divide by
+            # zero in the val-loss average.
+            self.val_dataloader = val if val is not None and len(val) > 0 else None
+            # The training loop calls set_epoch on self.sampler.
+            self.sampler = self.dataloader
+            self.val_sampler = None
+            return
+
         self.dataloader, self.sampler, self.val_dataloader, self.val_sampler = (
             self.setup_dataloader(
                 self.dset,
@@ -566,7 +607,10 @@ def _reconstruction_epoch(
         if inline_alignment:
             for ind in range(len(self.dset.tilt_angles)):
                 im_proj = proj_forward[:, ind, :]
-                im_meas = self.dset.forward(ind).target  # type: ignore
+                # proj_forward rows are volume slices (the tilt axis), i.e. the
+                # transpose of the stored tilt image -- the same orientation the
+                # error term below compares against.
+                im_meas = self.dset.forward(ind).target.T  # type: ignore
                 shift = torch_phase_cross_correlation(im_proj, im_meas)
                 if torch.linalg.norm(shift) <= 32:
                     shifted = torch.fft.ifft2(
@@ -585,7 +629,11 @@ def _reconstruction_epoch(
                         )
                     ).real
 
-                    proj_forward[:, ind, :] = shifted
+                    # Persist the aligned measurement in the tilt stack: the error
+                    # term reads the stack, and proj_forward is overwritten by
+                    # radon_torch below, so writing the aligned image there
+                    # silently discarded the alignment.
+                    self.dset.tilt_stack[ind] = shifted.T
 
         if mode == "sirt" or mode == "fbp":
             proj_forward = radon_torch(
diff --git a/src/quantem/tomography/tomography_base.py b/src/quantem/tomography/tomography_base.py
index 419b5ede..2301b7be 100644
--- a/src/quantem/tomography/tomography_base.py
+++ b/src/quantem/tomography/tomography_base.py
@@ -58,7 +58,11 @@ def __init__(
                 print("Setting up DDP for obj_model")
 
         self.dset = dset
-        self.dset.to(device)
+        # Use self.device (set by setup_distributed to cuda:local_rank under DDP),
+        # NOT the local `device` arg, which stays "cuda:0" on every rank and would
+        # strand the dataset's pose params on cuda:0 while the model/sampler live
+        # on cuda:local_rank (device-mismatch in _apply_pose's batched matmul).
+        self.dset.to(self.device)
 
     # --- Properties ---
     @property
diff --git a/src/quantem/tomography/utils.py b/src/quantem/tomography/utils.py
index 08093925..b1105d3c 100644
--- a/src/quantem/tomography/utils.py
+++ b/src/quantem/tomography/utils.py
@@ -5,10 +5,13 @@
 
 
 def rot_ZXZ(mags, z1, x, z3, device, mode="bilinear"):
-    if not isinstance(x, torch.Tensor) or not isinstance(z1, torch.Tensor):
-        z1 = torch.tensor(z1, dtype=torch.float32, device=device)
-        x = torch.tensor(x, dtype=torch.float32, device=device)
-        z3 = torch.tensor(z3, dtype=torch.float32, device=device)
+    # Convert each angle independently: re-wrapping an existing tensor with
+    # torch.tensor() copies and detaches it, silently cutting gradient flow
+    # through tensor angles whenever any other angle is passed as a float.
+    z1, x, z3 = (
+        a if isinstance(a, torch.Tensor) else torch.tensor(a, dtype=torch.float32, device=device)
+        for a in (z1, x, z3)
+    )
     curr_mags = mags
 
     curr_mags = differentiable_rotz_vectorized(curr_mags, z1, mode)
@@ -19,54 +22,54 @@ def rot_ZXZ(mags, z1, x, z3, device, mode="bilinear"):
     return curr_mags
 
 
-def differentiable_rotz_vectorized(mags, theta, mode="bilinear"):
-    _, dimz, dimy, dimx = mags.shape
-
-    if theta.dim() == 0:
-        theta = theta.unsqueeze(0)
-
+def _rot2d_affine_matrix(theta: torch.Tensor, batch: int) -> torch.Tensor:
+    """(B, 2, 3) in-plane rotation matrices, broadcasting a single angle over the batch."""
     theta_rad = torch.deg2rad(theta)
-
     cos_t, sin_t = torch.cos(theta_rad), torch.sin(theta_rad)
     affine_matrix = torch.stack(
         [cos_t, -sin_t, torch.zeros_like(theta), sin_t, cos_t, torch.zeros_like(theta)], dim=-1
     ).view(-1, 2, 3)
-
-    mags = mags.permute(1, 0, 2, 3)
-
-    def transform_slice(mag_slice):
-        grid = F.affine_grid(affine_matrix, mag_slice.unsqueeze(0).shape, align_corners=False)
-        return F.grid_sample(mag_slice.unsqueeze(0), grid, mode=mode, align_corners=False).squeeze(
-            0
+    if affine_matrix.shape[0] == 1 and batch > 1:
+        affine_matrix = affine_matrix.expand(batch, 2, 3)
+    elif affine_matrix.shape[0] != batch and batch != 1:
+        raise ValueError(
+            f"Got {affine_matrix.shape[0]} angles for a batch of {batch} volumes; "
+            "pass one angle, or one angle per volume."
         )
+    return affine_matrix
 
-    rotated_mags = torch.vmap(transform_slice)(mags)
-    return rotated_mags.permute(1, 0, 2, 3)
 
-
-def differentiable_rotx_vectorized(mags, theta, mode="bilinear"):
-    _, dimz, dimy, dimx = mags.shape
+def differentiable_rotz_vectorized(mags, theta, mode="bilinear"):
+    B, dimz, dimy, dimx = mags.shape
 
     if theta.dim() == 0:
         theta = theta.unsqueeze(0)
 
-    theta_rad = torch.deg2rad(theta)
+    # A rotation about z applies the same 2-D transform to every z-slice, so the
+    # z-axis rides along as grid_sample channels in a single call. (The previous
+    # per-slice vmap also broke for more than one angle, because affine_grid
+    # requires the matrix batch to match the slice batch of 1.)
+    affine_matrix = _rot2d_affine_matrix(theta, B)
+    if affine_matrix.shape[0] > B:  # one volume, many angles
+        mags = mags.expand(affine_matrix.shape[0], dimz, dimy, dimx)
+    grid = F.affine_grid(affine_matrix, mags.shape, align_corners=False)
+    return F.grid_sample(mags, grid, mode=mode, align_corners=False)
 
-    cos_t, sin_t = torch.cos(theta_rad), torch.sin(theta_rad)
-    affine_matrix = torch.stack(
-        [cos_t, -sin_t, torch.zeros_like(theta), sin_t, cos_t, torch.zeros_like(theta)], dim=-1
-    ).view(-1, 2, 3)
 
-    mags = mags.permute(3, 0, 1, 2)
+def differentiable_rotx_vectorized(mags, theta, mode="bilinear"):
+    B, dimz, dimy, dimx = mags.shape
 
-    def transform_slice(mag_slice):
-        grid = F.affine_grid(affine_matrix, mag_slice.unsqueeze(0).shape, align_corners=False)
-        return F.grid_sample(mag_slice.unsqueeze(0), grid, mode=mode, align_corners=False).squeeze(
-            0
-        )
+    if theta.dim() == 0:
+        theta = theta.unsqueeze(0)
 
-    rotated_mags = torch.vmap(transform_slice)(mags)
-    return rotated_mags.permute(1, 2, 3, 0)
+    # Same trick as rotz with the x-axis as the channel dim: rotate in (z, y).
+    affine_matrix = _rot2d_affine_matrix(theta, B)
+    mags = mags.permute(0, 3, 1, 2)  # (B, X, Z, Y)
+    if affine_matrix.shape[0] > B:  # one volume, many angles
+        mags = mags.expand(affine_matrix.shape[0], dimx, dimz, dimy)
+    grid = F.affine_grid(affine_matrix, mags.shape, align_corners=False)
+    rotated = F.grid_sample(mags, grid, mode=mode, align_corners=False)
+    return rotated.permute(0, 2, 3, 1)  # back to (B, Z, Y, X)
 
 
 def tv_loss_1d(x: torch.Tensor, reduction: str = "mean") -> torch.Tensor:
diff --git a/tests/ml/test_inr.py b/tests/ml/test_inr.py
new file mode 100644
index 00000000..c78bc1af
--- /dev/null
+++ b/tests/ml/test_inr.py
@@ -0,0 +1,40 @@
+"""Tests for ``quantem.core.ml.inr`` Siren / HSiren construction."""
+
+import torch
+
+from quantem.core.ml.inr import HSiren, Siren
+
+
+def _first_layer_weight(winner_seed) -> torch.Tensor:
+    # Fix the global RNG so the base SineLayer init is identical across builds;
+    # only the winner-initialization perturbation varies with winner_seed.
+    torch.manual_seed(0)
+    model = HSiren(hidden_layers=1, hidden_features=8, winner_initialization=winner_seed)
+    return model.net[0].linear.weight.detach().clone()
+
+
+class TestWinnerInitialization:
+    def test_same_seed_is_reproducible(self):
+        assert torch.equal(_first_layer_weight(72), _first_layer_weight(72))
+
+    def test_different_seeds_differ(self):
+        """Regression: the seeded torch.Generator was created but never used
+        (torch.randn_like ignores generators), so every winner seed produced
+        the same perturbation."""
+        assert not torch.equal(_first_layer_weight(72), _first_layer_weight(73))
+
+    def test_true_uses_default_seed(self):
+        assert torch.equal(_first_layer_weight(True), _first_layer_weight(42))
+
+    def test_disabled_adds_no_perturbation(self):
+        torch.manual_seed(0)
+        base = Siren(hidden_layers=1, hidden_features=8, winner_initialization=False)
+        torch.manual_seed(0)
+        again = Siren(hidden_layers=1, hidden_features=8, winner_initialization=False)
+        assert torch.equal(base.net[0].linear.weight.detach(), again.net[0].linear.weight.detach())
+
+
+def test_forward_shape():
+    model = HSiren(hidden_layers=1, hidden_features=8)
+    out = model(torch.rand(5, 3))
+    assert out.shape == (5, 1)
diff --git a/tests/ml/test_kplanes.py b/tests/ml/test_kplanes.py
new file mode 100644
index 00000000..bdd865bc
--- /dev/null
+++ b/tests/ml/test_kplanes.py
@@ -0,0 +1,20 @@
+"""Tests for ``quantem.core.ml.models.kplanes`` construction guards."""
+
+import pytest
+
+from quantem.core.ml.models.kplanes import KPlanes, KPlanesTILTED
+
+
+class TestResolutionValidation:
+    def test_isotropic_builds(self):
+        model = KPlanes(M_features=2, resolution=(16, 16, 16))
+        assert model.grids[0].shape == (3, 2, 16, 16)
+
+    def test_anisotropic_raises(self):
+        """Regression: the plane grids are allocated as (res[1], res[0]) for all three
+        axis pairs, ignoring res[2] -- anisotropic resolutions silently gave the
+        XZ/YZ planes the wrong grid along z instead of erroring."""
+        with pytest.raises(ValueError, match="isotropic"):
+            KPlanes(M_features=2, resolution=(16, 16, 8))
+        with pytest.raises(ValueError, match="isotropic"):
+            KPlanesTILTED(M_features=2, T=2, resolution=(16, 8, 16))
diff --git a/tests/ml/test_optimizermixin.py b/tests/ml/test_optimizermixin.py
index d6347ce6..7dc6f4e7 100644
--- a/tests/ml/test_optimizermixin.py
+++ b/tests/ml/test_optimizermixin.py
@@ -356,9 +356,7 @@ def test_parse_default_name_is_none(self):
 from quantem.core.ml.optimizer_mixin import OptimizerMixin  # noqa: E402
 
 torch = pytest.importorskip("torch") if config.get("has_torch") else None
-requires_torch = pytest.mark.skipif(
-    not config.get("has_torch"), reason="requires torch"
-)
+requires_torch = pytest.mark.skipif(not config.get("has_torch"), reason="requires torch")
 
 
 def _param(value=1.0):
@@ -413,7 +411,10 @@ def test_none_optimizer_removes_without_touching_params(self):
     def test_multi_group_pplr_applies_per_group_lr(self):
         model = _FakeModel({"descan": [_param()], "scan_positions": [_param(2.0)]})
         model.set_optimizer(
-            {"descan": OptimizerParams.SGD(lr=1e-2), "scan_positions": OptimizerParams.SGD(lr=1e-3)}
+            {
+                "descan": OptimizerParams.SGD(lr=1e-2),
+                "scan_positions": OptimizerParams.SGD(lr=1e-3),
+            }
         )
         groups = model.optimizer.param_groups
         assert len(groups) == 2
@@ -445,9 +446,56 @@ def test_reset_optimizer_on_unconfigured_model_is_noop(self):
     def test_set_scheduler_base_lr_uses_max_group_lr(self):
         model = _FakeModel({"descan": [_param()], "scan_positions": [_param(2.0)]})
         model.set_optimizer(
-            {"descan": OptimizerParams.SGD(lr=1e-2), "scan_positions": OptimizerParams.SGD(lr=1e-3)}
+            {
+                "descan": OptimizerParams.SGD(lr=1e-2),
+                "scan_positions": OptimizerParams.SGD(lr=1e-3),
+            }
         )
         model.set_scheduler(SchedulerParams.Plateau(), num_iter=10)
         assert model.scheduler is not None
         # Plateau min_lr defaults to base_LR / 20, with base_LR = max group lr (1e-2)
         assert model.scheduler.min_lrs[0] == pytest.approx(1e-2 / 20)
+
+
+class TestSchedulerParamsArePure:
+    """Regression: params() mutated the dataclass, so a shared instance baked in the
+    first call's derived values (min_lr, gamma, bounds, horizons) for every later
+    optimizer or reconstruct() call."""
+
+    def test_plateau_min_lr_follows_base_lr(self):
+        p = SchedulerParams.Plateau()
+        assert p.params(base_LR=1.0)["min_lr"] == pytest.approx(1 / 20)
+        assert p.params(base_LR=10.0)["min_lr"] == pytest.approx(10 / 20)
+        assert p.min_lr is None  # instance untouched
+
+    def test_exponential_gamma_follows_num_iter(self):
+        p = SchedulerParams.Exponential(factor=0.5)
+        g10 = p.params(base_LR=1.0, num_iter=10)["gamma"]
+        g100 = p.params(base_LR=1.0, num_iter=100)["gamma"]
+        assert g10 == pytest.approx(0.5 ** (1 / 10))
+        assert g100 == pytest.approx(0.5 ** (1 / 100))
+        assert p.num_iter is None and p.gamma == 0.9
+
+    def test_cyclic_bounds_follow_base_lr(self):
+        p = SchedulerParams.Cyclic()
+        first = p.params(base_LR=1.0)
+        second = p.params(base_LR=2.0)
+        assert second["base_lr"] == pytest.approx(2 * first["base_lr"])
+        assert second["max_lr"] == pytest.approx(2 * first["max_lr"])
+        assert p.base_lr is None and p.max_lr is None
+
+    def test_linear_and_cosine_follow_num_iter(self):
+        lin = SchedulerParams.Linear()
+        assert lin.params(base_LR=1.0, num_iter=5)["total_iters"] == 5
+        assert lin.params(base_LR=1.0, num_iter=50)["total_iters"] == 50
+        assert lin.total_iters is None
+        cos = SchedulerParams.CosineAnnealing()
+        assert cos.params(base_LR=1.0, num_iter=5)["T_max"] == 5
+        assert cos.params(base_LR=1.0, num_iter=50)["T_max"] == 50
+        assert cos.T_max is None
+
+    def test_explicit_values_still_win(self):
+        p = SchedulerParams.Plateau(min_lr=1e-6)
+        assert p.params(base_LR=1.0)["min_lr"] == 1e-6
+        e = SchedulerParams.Exponential(gamma=0.8)
+        assert e.params(base_LR=1.0, num_iter=10)["gamma"] == 0.8
diff --git a/tests/tomography/test_dataset_models.py b/tests/tomography/test_dataset_models.py
index 04abd3b5..268fa2c6 100644
--- a/tests/tomography/test_dataset_models.py
+++ b/tests/tomography/test_dataset_models.py
@@ -58,6 +58,20 @@ def test_normalised_by_95th_quantile(self):
         q95 = torch.quantile(d.tilt_stack, 0.95)
         assert torch.isclose(q95, torch.tensor(1.0), atol=1e-4)
 
+    def test_sparse_stack_normalisation_stays_finite(self):
+        # >95% zeros makes the 95th quantile 0; normalisation must not divide
+        # by it (inf/NaN targets poison every parameter on the first backward).
+        stack = np.zeros((5, 12, 12), dtype=np.float32)
+        stack[:, 5:7, 5:7] = 10.0
+        d = TomographyPixDataset.from_data(stack, np.linspace(-60, 60, 5).astype(np.float32))
+        assert torch.isfinite(d.tilt_stack).all()
+        assert d.tilt_stack.max() > 0
+
+    def test_all_zero_stack_raises(self):
+        stack = np.zeros((5, 12, 12), dtype=np.float32)
+        with pytest.raises(ValueError):
+            TomographyPixDataset.from_data(stack, np.linspace(-60, 60, 5).astype(np.float32))
+
     def test_reference_idx_and_learnable_tilts(self):
         # negated angles -> [40, 15, -10, -35, -60]; smallest |angle| is index 2.
         angles = np.linspace(-40, 60, 5).astype(np.float32)
@@ -180,6 +194,25 @@ def test_getitem_index_mapping(self):
         assert int(item["pixel_j"]) == 1
         assert torch.isclose(item["target_value"], d.tilt_stack[1, 1, 1])
 
+    def test_len_and_getitem_non_square(self):
+        # Regression: pixel_i/pixel_j must decompose by the width (shape[2]) and __len__
+        # by H*W, not max(shape)^2 -- both were wrong for rectangular tilt images.
+        rng = np.random.default_rng(0)
+        stack = rng.random((3, 4, 6)).astype(np.float32)  # nang=3, H=4, W=6
+        d = TomographyINRDataset.from_data(stack, np.linspace(-60, 60, 3, dtype="f4"))
+        assert len(d) == 3 * 4 * 6
+        # idx = proj*(H*W) + i*W + j = 1*24 + 2*6 + 3 = 39
+        item = d[39]
+        assert int(item["projection_idx"]) == 1
+        assert int(item["pixel_i"]) == 2
+        assert int(item["pixel_j"]) == 3
+        assert torch.isclose(item["target_value"], d.tilt_stack[1, 2, 3])
+        # every index in range maps to a valid pixel
+        last = d[len(d) - 1]
+        assert int(last["projection_idx"]) == 2
+        assert int(last["pixel_i"]) == 3
+        assert int(last["pixel_j"]) == 5
+
     def test_save_load_parameters_roundtrip(self, tmp_path):
         angles = np.linspace(-60, 60, 5, dtype="f4")
         d = TomographyINRDataset.from_data(_stack(), angles)
@@ -192,3 +225,30 @@ def test_save_load_parameters_roundtrip(self, tmp_path):
         d2.to("cpu")
         d2.load_parameters(path)
         assert torch.allclose(d2.z1_params.detach(), d.z1_params.detach())
+
+
+@pytest.mark.parametrize("cls", [TomographyPixDataset, TomographyINRDataset])
+def test_to_preserves_trained_pose_parameters(cls):
+    """Regression: to() rebuilt the pose parameters from the initial-value buffers,
+    so any device move after training (e.g. from_file(...).to(device)) silently
+    reset the learned pose to zero."""
+    d = cls.from_data(_stack(), np.linspace(-60, 60, 5, dtype="f4"))
+    d.to("cpu")
+    d.z1_params.data.fill_(0.37)
+    d.z3_params.data.fill_(-0.21)
+    d.shifts_params.data.fill_(1.5)
+
+    d.to("cpu")
+
+    assert torch.allclose(d.z1_params.detach(), torch.full_like(d.z1_params, 0.37))
+    assert torch.allclose(d.z3_params.detach(), torch.full_like(d.z3_params, -0.21))
+    assert torch.allclose(d.shifts_params.detach(), torch.full_like(d.shifts_params, 1.5))
+
+
+def test_learnable_tilts_is_read_only():
+    """Regression: the old setter wrote a private attribute the getter never read,
+    so assignment appeared to succeed while silently doing nothing."""
+    d = TomographyPixDataset.from_data(_stack(), np.linspace(-60, 60, 5).astype(np.float32))
+    with pytest.raises(AttributeError):
+        d.learnable_tilts = 3
+    assert d.learnable_tilts == 4
diff --git a/tests/tomography/test_device_batch_sampler.py b/tests/tomography/test_device_batch_sampler.py
new file mode 100644
index 00000000..8e050752
--- /dev/null
+++ b/tests/tomography/test_device_batch_sampler.py
@@ -0,0 +1,147 @@
+"""Tests for ``DeviceBatchSampler`` and its wiring in ``Tomography.reconstruct``.
+
+The sampler must yield batches identical in content to the per-pixel
+DataLoader path (same keys, same index decode as
+``TomographyINRDataset.__getitem__``), cover each pixel exactly once per
+epoch (minus the dropped tail batch), and respect the train/val index
+split. Device-independent, so everything here runs on CPU.
+"""
+
+import numpy as np
+import torch
+
+from quantem.tomography.dataset_models import DeviceBatchSampler, TomographyINRDataset
+
+
+def _dset(n_proj=4, n=10, seed=0):
+    rng = np.random.default_rng(seed)
+    stack = rng.random((n_proj, n, n)).astype(np.float32)
+    angles = np.linspace(-60, 60, n_proj).astype(np.float32)
+    return TomographyINRDataset.from_data(tilt_stack=stack, tilt_angles=angles)
+
+
+def test_batches_match_getitem():
+    dset = _dset()
+    sampler = DeviceBatchSampler(dset, batch_size=37, device="cpu", shuffle=False)
+    seen = 0
+    for batch in sampler:
+        for k in range(len(batch["target_value"])):
+            item = dset[seen + k]
+            # __getitem__ returns plain ints for the index keys (cheaper than
+            # 0-d tensors); coerce before comparing against the batched tensors.
+            for key in ("projection_idx", "pixel_i", "pixel_j", "phi", "target_value"):
+                torch.testing.assert_close(
+                    batch[key][k],
+                    torch.as_tensor(item[key], dtype=batch[key].dtype),
+                    rtol=0,
+                    atol=0,
+                )
+        seen += len(batch["target_value"])
+
+
+def test_epoch_covers_indices_once_with_drop_last():
+    dset = _dset()
+    n = len(dset)
+    batch_size = 64
+    sampler = DeviceBatchSampler(dset, batch_size=batch_size, device="cpu", shuffle=True)
+    assert len(sampler) == n // batch_size
+    per_proj = dset.tilt_stack.shape[1] * dset.tilt_stack.shape[2]
+    flat = []
+    for batch in sampler:
+        assert len(batch["target_value"]) == batch_size
+        flat.append(
+            batch["projection_idx"] * per_proj
+            + batch["pixel_i"] * dset.tilt_stack.shape[1]
+            + batch["pixel_j"]
+        )
+    flat = torch.cat(flat)
+    assert flat.unique().numel() == flat.numel()  # no repeats within an epoch
+    assert flat.numel() == len(sampler) * batch_size
+
+
+def test_shuffle_changes_order_between_epochs():
+    dset = _dset()
+    sampler = DeviceBatchSampler(dset, batch_size=50, device="cpu", shuffle=True)
+    first = next(iter(sampler))["target_value"]
+    second = next(iter(sampler))["target_value"]
+    assert not torch.equal(first, second)
+
+
+def test_val_split_is_disjoint():
+    dset = _dset()
+    n = len(dset)
+    perm = torch.randperm(n)
+    n_val = n // 10
+    train = DeviceBatchSampler(dset, 32, "cpu", indices=perm[n_val:])
+    val = DeviceBatchSampler(dset, 32, "cpu", indices=perm[:n_val], shuffle=False)
+    assert len(train._indices) + len(val._indices) == n
+    assert torch.cat([train._indices, val._indices]).unique().numel() == n
+
+
+def test_reconstruct_uses_sampler_single_process():
+    """Smoke: the single-process reconstruct path builds DeviceBatchSamplers.
+
+    Object models that go through ``setup_distributed`` must be built on CUDA
+    when a CUDA device exists, so pick the device accordingly.
+    """
+    from quantem.core.ml.models.kplanes import KPlanesTILTED
+    from quantem.core.ml.optimizer_mixin import OptimizerParams
+    from quantem.tomography.object_models import ObjectINR
+    from quantem.tomography.tomography import Tomography
+
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    dset = _dset(n_proj=3, n=8)
+    model = KPlanesTILTED(M_features=2, resolution=(8, 8, 8), multiscale_res_multipliers=[1], T=1)
+    obj = ObjectINR.from_model(model, shape=(8, 8, 8), device=device)
+    tomo = Tomography.from_models(dset=dset, obj_model=obj, device=device, verbose=False)
+    tomo.reconstruct(
+        num_iter=1,
+        batch_size=32,
+        num_workers=0,
+        val_fraction=0.25,
+        optimizer_params={
+            "object": {"default": OptimizerParams.Adam(lr=1e-3)},
+            "pose": OptimizerParams.Adam(lr=1e-3),
+        },
+    )
+    assert isinstance(tomo.dataloader, DeviceBatchSampler)
+    assert isinstance(tomo.val_dataloader, DeviceBatchSampler)
+    # the loop drives epoch reshuffling through set_epoch on self.sampler
+    assert tomo.sampler is tomo.dataloader
+
+
+def _flat(batch, s1, s2):
+    return batch["projection_idx"] * (s1 * s2) + batch["pixel_i"] * s1 + batch["pixel_j"]
+
+
+def test_ddp_shards_are_disjoint_and_equal():
+    dset = _dset()
+    s1, s2 = dset.tilt_stack.shape[1], dset.tilt_stack.shape[2]
+    world = 4
+    shards = []
+    for rank in range(world):
+        sampler = DeviceBatchSampler(dset, 32, "cpu", rank=rank, world_size=world)
+        sampler.set_epoch(3)
+        flats = torch.cat([_flat(b, s1, s2) for b in sampler])
+        shards.append(flats)
+        assert len(sampler) == (len(dset) // world) // 32  # equal on every rank
+    allv = torch.cat(shards)
+    assert allv.unique().numel() == allv.numel()  # no pixel on two ranks
+
+
+def test_ddp_epoch_permutation_shared_and_reproducible():
+    dset = _dset()
+    s1, s2 = dset.tilt_stack.shape[1], dset.tilt_stack.shape[2]
+
+    def epoch_flats(rank, epoch):
+        sampler = DeviceBatchSampler(dset, 32, "cpu", rank=rank, world_size=2)
+        sampler.set_epoch(epoch)
+        return torch.cat([_flat(b, s1, s2) for b in sampler])
+
+    # same (rank, epoch) on a fresh instance -> identical batches
+    torch.testing.assert_close(epoch_flats(0, 7), epoch_flats(0, 7), rtol=0, atol=0)
+    # different epoch -> different order
+    assert not torch.equal(epoch_flats(0, 7), epoch_flats(0, 8))
+    # ranks of the same epoch are disjoint
+    both = torch.cat([epoch_flats(0, 7), epoch_flats(1, 7)])
+    assert both.unique().numel() == both.numel()
diff --git a/tests/tomography/test_object_models.py b/tests/tomography/test_object_models.py
index c4216827..da01c7d8 100644
--- a/tests/tomography/test_object_models.py
+++ b/tests/tomography/test_object_models.py
@@ -148,6 +148,17 @@ def test_soft_constraint_zero_when_tv_off(self):
         obj.constraints.tv_vol = 0.0
         assert float(obj.apply_soft_constraints(ctx).detach()) == 0.0
 
+    def test_soft_constraint_with_tv_on_backprops(self):
+        # Regression: soft_loss was a leaf tensor with requires_grad=True, so the
+        # in-place `+= tv_loss` raised RuntimeError whenever tv_vol > 0.
+        ctx = ReconstructionContext(obj=torch.rand(8, 8, 8).requires_grad_(True))
+        obj = ObjectPixelated.from_uniform(shape=(8, 8, 8), device="cpu")
+        obj.constraints.tv_vol = 1.0
+        loss = obj.apply_soft_constraints(ctx)
+        assert torch.isfinite(loss) and loss > 0.0
+        loss.backward()
+        assert ctx.obj.grad is not None
+
 
 class TestFactoryGuard:
     def test_objectbase_requires_token(self):
@@ -205,10 +216,25 @@ def _obj(self, device, n=16):
 
     def test_forward_masks_out_of_range(self, torch_device):
         obj = self._obj(torch_device)
-        coords = torch.tensor([[0.0, 0.0, 0.0], [5.0, 0.0, 0.0]], device=torch_device)
+        # Regression: the mask used to check x and y only, so tilted rays leaving the
+        # volume along z still contributed (extrapolated) density to the integral.
+        coords = torch.tensor(
+            [[0.0, 0.0, 0.0], [5.0, 0.0, 0.0], [0.0, 5.0, 0.0], [0.0, 0.0, 5.0]],
+            device=torch_device,
+        )
         out = obj.forward(coords)
-        assert out.shape[0] == 2
-        assert float(out[1].detach()) == 0.0  # x=5 is outside [-1, 1] -> masked to zero
+        assert out.shape[0] == 4
+        for i in (1, 2, 3):  # each axis outside [-1, 1] -> masked to zero
+            assert float(out[i].detach()) == 0.0
+
+    def test_soft_constraints_without_coords(self, torch_device):
+        # Regression: soft_loss was created on ctx.coords.device before the None
+        # check, raising AttributeError when no constraints are active.
+        obj = self._obj(torch_device)
+        obj.constraints.tv_vol = 0.0
+        obj.constraints.sparsity = 0.0
+        loss = obj.apply_soft_constraints(ReconstructionContext())
+        assert float(loss.detach()) == 0.0
 
     def test_apply_hard_constraints_positivity(self, torch_device):
         obj = self._obj(torch_device)
@@ -257,6 +283,17 @@ def test_volume_tv_loss_scalar(self, torch_device):
         assert loss.ndim == 0
         assert torch.isfinite(loss)
 
+    def test_soft_constraints_plane_tv_only(self, torch_device):
+        """Regression: tv_plane > 0 with tv_vol = 0 must still apply the plane TV.
+        The gate previously keyed on tv_vol alone, silently dropping plane-only TV."""
+        obj = self._obj(torch_device)
+        obj.constraints.tv_plane = 0.1
+        obj.constraints.tv_vol = 0.0
+        coords = torch.rand(64, 3, device=torch_device) * 2 - 1
+        ctx = ReconstructionContext(coords=coords, pred=torch.zeros(64, device=torch_device))
+        loss = obj.apply_soft_constraints(ctx)
+        assert float(loss.detach()) > 0.0
+
     def test_normalize_optimizer_params_rejects_non_dict(self, torch_device):
         from quantem.core.ml.optimizer_mixin import OptimizerParams
 
diff --git a/tests/tomography/test_tomography_conventional.py b/tests/tomography/test_tomography_conventional.py
index a7cfb5e0..85e3daf2 100644
--- a/tests/tomography/test_tomography_conventional.py
+++ b/tests/tomography/test_tomography_conventional.py
@@ -61,6 +61,27 @@ def test_obj_constraints_accepts_dict(self, phantom_volume, tilt_series, tilt_an
         assert tomo.num_epochs == 2
 
 
+class TestInlineAlignment:
+    def test_alignment_corrects_misaligned_projection(
+        self, phantom_volume, tilt_series, tilt_angles
+    ):
+        """Regression: the aligned measurement was written into proj_forward, which
+        radon_torch overwrites immediately, so inline_alignment was a no-op."""
+        n = phantom_volume.shape[0]
+        shifted_series = tilt_series.copy()
+        shifted_series[4] = np.roll(shifted_series[4], shift=3, axis=1)
+
+        reference = _build(tilt_series, tilt_angles, n).dset.tilt_stack[4]
+        tomo = _build(shifted_series, tilt_angles, n)
+        misaligned_before = tomo.dset.tilt_stack[4].clone()
+        tomo.reconstruct(num_iter=3, mode="sirt", inline_alignment=True)
+        aligned_after = tomo.dset.tilt_stack[4]
+
+        err_before = float(((misaligned_before - reference) ** 2).mean())
+        err_after = float(((aligned_after - reference) ** 2).mean())
+        assert err_after < err_before  # the stack was actually re-aligned
+
+
 class TestFBP:
     def test_fbp_runs_single_epoch(self, phantom_volume, tilt_series, tilt_angles):
         n = phantom_volume.shape[0]
diff --git a/tests/tomography/test_utils.py b/tests/tomography/test_utils.py
index 1fa6a47c..b8530491 100644
--- a/tests/tomography/test_utils.py
+++ b/tests/tomography/test_utils.py
@@ -5,6 +5,7 @@
 import torch
 
 from quantem.tomography.utils import (
+    differentiable_rotx_vectorized,
     differentiable_rotz_vectorized,
     rot_ZXZ,
     tv_loss_1d,
@@ -76,3 +77,37 @@ def test_gradient_flows_through_rotation(self):
         out.sum().backward()
         assert vol.grad is not None
         assert torch.isfinite(vol.grad).all()
+
+
+class TestRotZXZGradients:
+    def test_grad_flows_with_mixed_float_and_tensor_angles(self):
+        """Regression: a non-tensor angle made rot_ZXZ re-wrap every angle with
+        torch.tensor(), detaching gradients through tensor angles."""
+        vol = _block_volume()
+        x = torch.tensor(20.0, requires_grad=True)
+        out = rot_ZXZ(vol, 0.0, x, 0.0, device="cpu")
+        out.sum().backward()
+        assert x.grad is not None
+        assert torch.isfinite(x.grad)
+    @pytest.mark.parametrize(
+        "rot_fn", [differentiable_rotz_vectorized, differentiable_rotx_vectorized]
+    )
+    def test_multi_angle_matches_per_angle_calls(self, rot_fn):
+        # Regression: the per-slice vmap version raised for more than one angle
+        # (affine_grid requires the matrix batch to match the slice batch of 1).
+        vol = _block_volume()
+        angles = torch.tensor([10.0, -25.0, 60.0])
+        multi = rot_fn(vol, angles)
+        per_angle = torch.cat([rot_fn(vol, a) for a in angles])
+        assert multi.shape == (3, *vol.shape[1:])
+        assert torch.allclose(multi, per_angle, atol=1e-6)
+
+    @pytest.mark.parametrize(
+        "rot_fn", [differentiable_rotz_vectorized, differentiable_rotx_vectorized]
+    )
+    def test_per_volume_angles(self, rot_fn):
+        vols = torch.cat([_block_volume(), _block_volume().flip(-1), _block_volume().flip(-2)])
+        angles = torch.tensor([10.0, -25.0, 60.0])
+        batched = rot_fn(vols, angles)
+        per_volume = torch.cat([rot_fn(vols[i : i + 1], angles[i]) for i in range(3)])
+        assert torch.allclose(batched, per_volume, atol=1e-6)