start move over to previous engine generation code for upscale

TNTwise · Oct 27, 2024 · a694e44 · a694e44
1 parent 4bcca51
commit a694e44
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/backend/src/spandrel/architectures/SPAN/__arch/span.py b/backend/src/spandrel/architectures/SPAN/__arch/span.py
@@ -284,7 +284,10 @@ def is_norm(self):
         return self.no_norm is None
 
     def forward(self, x):
+        dtype = x.dtype
+        device = x.device
         if self.is_norm:
+            self.mean.to(device=device, dtype=dtype)
             self.mean = self.mean.type_as(x)
             x = (x - self.mean) * self.img_range