Update app.py

app.py

@@ -29,7 +29,7 @@ sampler_list = [
 torch.backends.cudnn.deterministic = True
 torch.backends.cudnn.benchmark = False
 
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
@@ -72,7 +72,7 @@ def load_pipeline(model_name):
             add_watermarker=False,
             use_auth_token=HF_TOKEN
         )
-    pipe.to("cuda:0")
+    pipe.to(device)
     return pipe
 
 @spaces.GPU
@@ -118,6 +118,7 @@ def generate(
 
 if torch.cuda.is_available():
     pipe = load_pipeline(MODEL)
+    pipe.to(device)
     print("Loaded on Device!")
 else:
     pipe = None