Spaces:

SLAYTECH
/

Dress_model

Build error

uhdessai commited on May 6, 2025

Commit

31f241b

verified ·

1 Parent(s): 36ac0c4

Update gen_images.py

Files changed (1) hide show

gen_images.py CHANGED Viewed

@@ -103,7 +103,7 @@ def generate_images(
     """
     print('Loading networks from "%s"...' % network_pkl)
-    device = torch.device('cuda')
     with dnnlib.util.open_url(network_pkl) as f:
         G = legacy.load_network_pkl(f)['G_ema'].to(device) # type: ignore

     """
     print('Loading networks from "%s"...' % network_pkl)
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     with dnnlib.util.open_url(network_pkl) as f:
         G = legacy.load_network_pkl(f)['G_ema'].to(device) # type: ignore