import spaces

app.py

@@ -9,13 +9,15 @@ from torchvision.utils import save_image #to save the generated images
 from tqdm import tqdm # progress bar
 import matplotlib.pyplot as plt
 import gradio as gr
+import spaces
 
 from styleTransfer import style_transfer
 from dataTransform import tensor_to_image
 
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+device = 'cuda' 
 print(device)
 
+@spaces.GPU
 def gradio_style_transfer(steps, content_image, style_image):
     generated_tensor = style_transfer(content_image, style_image, total_steps= steps)
     generated_image = tensor_to_image(generated_tensor)