From a29060424dc4c841dbdb100aeeca07178aba77c7 Mon Sep 17 00:00:00 2001
From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com>
Date: Tue, 27 Aug 2024 19:18:27 -0400
Subject: [PATCH] get torch device

---
 nodes/masks/flex_masks.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/nodes/masks/flex_masks.py b/nodes/masks/flex_masks.py
index fbbe4b5..1d02ad9 100644
--- a/nodes/masks/flex_masks.py
+++ b/nodes/masks/flex_masks.py
@@ -3,6 +3,7 @@ from .mask_utils import morph_mask, warp_mask, transform_mask, combine_masks
 import math
 import numpy as np
 from .voronoi_noise import VoronoiNoise #NOTE credit for Voronoi goes to Alan Huang https://github.com/alanhuang67/
+from comfy.model_management import get_torch_device
 
 class FlexMaskMorph(FlexMaskBase):
     @classmethod
@@ -184,7 +185,7 @@ class FlexMaskVoronoiScheduled(FlexMaskBase):
             Y=[y_offset],
             distance_metric=distance_metric,
             batch_size=1,
-            device='cpu'
+            device=get_torch_device()
         )
 
         # Generate Voronoi noise