Upload app.py

app.py

@@ -21,7 +21,7 @@ def _get_layers(arch:str, pretrained=True):
   "Get the layers and arch for a VGG Model (16 and 19 are supported only)"
   feat_net = vgg19(pretrained=pretrained) if arch.find('9') > 1 else vgg16(pretrained=pretrained)
   config = _vgg_config.get(arch)
-  features = feat_net.features.cuda().eval()
+  features = feat_net.features.eval()
   for p in features.parameters(): p.requires_grad=False
   return feat_net, [features[i] for i in config]
 
@@ -38,7 +38,7 @@ learner = from_pretrained_fastai(repo_id)
 
 def infer(img):
     pred = learner.predict(img)
-    image = pred[0].cpu().numpy()
+    image = pred[0].numpy()
     image = image.transpose((1, 2, 0))
     plt.imshow(image)
     return plt.gcf() #pred[0].show()