From b9db8bef89e8016843f92d12089ce62951e1231f Mon Sep 17 00:00:00 2001
From: Stephan <31624652+stephandooper@users.noreply.github.com>
Date: Wed, 12 Aug 2026 16:06:59 +0200
Subject: [PATCH] Preserve CPU statistics configuration

---
 lightstream/core/scnn/scnn.py   | 16 ++++++++--------
 lightstream/modules/__init__.py | 17 +++++++++++++----
 2 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/lightstream/core/scnn/scnn.py b/lightstream/core/scnn/scnn.py
index 822ad925..c66afc1f 100644
--- a/lightstream/core/scnn/scnn.py
+++ b/lightstream/core/scnn/scnn.py
@@ -8,7 +8,6 @@
 from dataclasses import dataclass
 from typing import List
 
-import numpy as np
 import torch
 import torch.autograd
 import torch.backends
@@ -331,6 +330,7 @@ def _slice_reducer_mask(
 
     def _configure(self):
         # Save current model and cudnn flags, since we need to change them and restore later
+        original_device = self.device
         state_dict = self._save_parameters()
         (
             old_deterministic_flag,
@@ -362,8 +362,8 @@ def _configure(self):
 
         # TODO; temp hack for tile sizes too big on gpu,
         if self.statistics_on_cpu:
-            self.stream_module = self.stream_module.cuda()
-            self.device = torch.device("cuda")  # type:ignore
+            self.stream_module = self.stream_module.to(original_device)
+            self.device = original_device
 
         # Remove all hooks and add hooks for correcting gradients
         # during lightstream
@@ -2335,15 +2335,15 @@ def _backward_gather_statistics_hook(self, module, grad_in, grad_out):
 
                 f_grad = torch.sum(grad_out[0], dim=1)[0]
                 f_grad = f_grad * new_outpt
-                f_grad = f_grad.cpu()
-                f_grad = np.repeat(f_grad, stride[1], axis=0)
-                f_grad = np.repeat(f_grad, stride[2], axis=1)
-                grad = np.zeros(grad_in[0].shape[2:])
+                f_grad = torch.repeat_interleave(f_grad, int(stride[1]), dim=0)
+                f_grad = torch.repeat_interleave(f_grad, int(stride[2]), dim=1)
+                grad = torch.zeros(
+                    grad_in[0].shape[2:], device=f_grad.device, dtype=f_grad.dtype
+                )
 
                 self._print_verbose("testing shape gradient fix")
                 grad[: f_grad.shape[0], : f_grad.shape[1]] = f_grad[: grad.shape[0], : grad.shape[1]]
 
-                f_grad = torch.from_numpy(grad)
                 f_grad = f_grad.to(self.device)
 
             if grad_out[0].numel() == 0 or torch.count_nonzero(grad_out[0]) == 0:
diff --git a/lightstream/modules/__init__.py b/lightstream/modules/__init__.py
index 5414f749..b268d4dd 100644
--- a/lightstream/modules/__init__.py
+++ b/lightstream/modules/__init__.py
@@ -1,8 +1,17 @@
-from .lightningstreaming import LightningStreamingModule
-from .imagenet_template import ImageNetClassifier
-
-
 __all__ = [
     "LightningStreamingModule",
     "ImageNetClassifier",
 ]
+
+
+def __getattr__(name):
+    """Load Lightning-dependent convenience classes only when requested."""
+    if name == "LightningStreamingModule":
+        from .lightningstreaming import LightningStreamingModule
+
+        return LightningStreamingModule
+    if name == "ImageNetClassifier":
+        from .imagenet_template import ImageNetClassifier
+
+        return ImageNetClassifier
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")