Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import torchvision
+import torchvision.transforms as transforms
 
 def normalize():
   MEAN = [0.485, 0.456, 0.406]
@@ -35,7 +36,7 @@ def tensor_to_img(tensor):
 def style_transfer(content_img, style_strength, style_img_1 = None, iw_1 = 0., style_img_2 = None, iw_2 = 0., style_img_3 = None, iw_3 = 0., preserve_color = None):
     transform = transformer(imsize = 512)
 
-    content = transform(content_img).unsqueeze(0).cuda()
+    content = transform(content_img).unsqueeze(0)
 
     iw = [iw_1, iw_2, iw_3]
     interpolation_weights = [i/ sum(iw) for i in iw]
@@ -44,7 +45,7 @@ def style_transfer(content_img, style_strength, style_img_1 = None, iw_1 = 0., s
     styles = []
     for style_img in style_imgs:
       if style_img is not None:
-        styles.append(transform(style_img).unsqueeze(0).cuda())
+        styles.append(transform(style_img).unsqueeze(0))
     if preserve_color == "None": preserve_color = None
     elif preserve_color == "Whitening": preserve_color = "batch_wct"
     #elif preserve_color == "Histogram matching": preserve_color = "histogram_matching"
@@ -80,4 +81,4 @@ interface = gr.Interface(fn = style_transfer,
                          description = None
                          )
 interface.queue()
-interface.launch(share = True, debug = True)
+interface.launch(share = True)