From 111d674e71f9ed01574471a0a511a36a07d72822 Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Thu, 6 Apr 2023 18:35:18 +1000 Subject: [PATCH] fix(nodes): use correct torch device in NoiseInvocation --- invokeai/app/invocations/latent.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/invokeai/app/invocations/latent.py b/invokeai/app/invocations/latent.py index 0481282ba9..49c3c4f11e 100644 --- a/invokeai/app/invocations/latent.py +++ b/invokeai/app/invocations/latent.py @@ -6,7 +6,7 @@ from torch import Tensor import torch from ...backend.model_management.model_manager import ModelManager -from ...backend.util.devices import CUDA_DEVICE, torch_dtype +from ...backend.util.devices import choose_torch_device, torch_dtype from ...backend.stable_diffusion.diffusion.shared_invokeai_diffusion import PostprocessingSettings from ...backend.image_util.seamless import configure_model_padding from ...backend.prompting.conditioning import get_uc_and_c_and_ec @@ -110,7 +110,7 @@ class NoiseInvocation(BaseInvocation): height: int = Field(default=512, multiple_of=64, gt=0, description="The height of the resulting noise", ) def invoke(self, context: InvocationContext) -> NoiseOutput: - device = torch.device(CUDA_DEVICE) + device = torch.device(choose_torch_device()) noise = get_noise(self.width, self.height, device, self.seed) name = f'{context.graph_execution_state_id}__{self.id}'