Update app.py

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 from train import main
 import spaces
-import torch
 import tensorflow as tf
 
 
@@ -40,9 +39,8 @@ def handle_image_size_selection(img_size, content_img):
 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")
     print(tf.config.list_physical_devices('GPU'))
-    print(torch.cuda.get_device_name(0))
-    with tf.device("/physical_device:GPU: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)
+    # with tf.device("/physical_device:GPU: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
 
 # @spaces.GPU