Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from train import main
 import spaces
+import tensorflow as tf
 
 # Function to calculate the aspect ratio
 def calculate_aspect_ratio(width, height):
@@ -36,7 +37,7 @@ def handle_image_size_selection(img_size, content_img):
 @spaces.GPU
 def process_images(content_img, style_img, epochs, steps_per_epoch, learning_rate, content_loss_factor, style_loss_factor, img_size, img_width, img_height):
     print("Start processing")
-    with device("/cuda:0"):
+    with tf.device("/cuda:0"):
         output_img = main(content_img, style_img, epochs, steps_per_epoch, learning_rate, content_loss_factor, style_loss_factor, img_size, img_width, img_height)
     return output_img