Update app.py

app.py

@@ -16,22 +16,22 @@ from model_video import build_model
 import numpy as np
 import collections
 #import argparse
-
-net = build_model('cpu').to('cpu')
+device='cuda:0'
+net = build_model(device).to(device)
     #net=torch.nn.DataParallel(net)
 model_path = 'image_best.pth'
 print(model_path)
-weight=torch.load(model_path,map_location=torch.device('cpu'))
+weight=torch.load(model_path,map_location=torch.device(device))
 #print(type(weight))
 new_dict=collections.OrderedDict()
 for k in weight.keys():
   new_dict[k[len('module.'):]]=weight[k]
 net.load_state_dict(new_dict)
 net.eval()
-net = net.to('cpu')
+net = net.to(device)
 def test(gpu_id, net, img_list, group_size, img_size):
     print('test')
-    device='cpu'
+    #device=device
     
     img_transform = transforms.Compose([transforms.Resize((img_size, img_size)), transforms.ToTensor(),
                     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
@@ -65,7 +65,7 @@ def sepia(img1,img2,img3,img4,img5):
   h_list,w_list=[_.shape[0] for _ in img_list],[_.shape[1] for _ in img_list]
   #print(type(img1))
   #print(img1.shape)
-  result_list=test('cpu',net,img_list,5,224)
+  result_list=test(device,net,img_list,5,224)
   #result_list=[result_list[i].resize((w_list[i], h_list[i]), Image.BILINEAR) for i in range(5)]
   img1,img2,img3,img4,img5=result_list#test('cpu',net,img_list,5,224)
   return img1,img2,img3,img4,img5