Only do .to if necessary (#2859)

only do .to if necessary
chaiNNer-org · May 17, 2024 · 268e683 · 268e683
1 parent a52151a
commit 268e683
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/backend/src/nodes/impl/pytorch/auto_split.py b/backend/src/nodes/impl/pytorch/auto_split.py
@@ -58,7 +58,8 @@ def pytorch_auto_split(
     progress: Progress,
 ) -> np.ndarray:
     dtype = torch.float16 if use_fp16 else torch.float32
-    model = model.to(device, dtype)
+    if model.dtype != dtype or model.device != device:
+        model = model.to(device, dtype)
 
     def upscale(img: np.ndarray, _: object):
         progress.check_aborted()