update

app.py

@@ -9,11 +9,11 @@ import torchvision.transforms.functional as TF
 channels = [64, 128, 256, 512]
 layers = [2, 2, 4, 2]
 num_subnet = 4
-net_i = FullNet_NLP(channels, layers, num_subnet, 4,num_classes=1000, drop_path=0,save_memory=True, inter_supv=True, head_init_scale=None,kernel_size=3).cpu()
+net_i = FullNet_NLP(channels, layers, num_subnet, 4,num_classes=1000, drop_path=0,save_memory=True, inter_supv=True, head_init_scale=None,kernel_size=3).cuda()
 for param in net_i.parameters():
     param.data = param.data.to(torch.float16)
 net_i.load_state_dict(torch.load('./fp16_check.pt')['icnn'])
-net_c = PretrainedConvNext("convnext_small_in22k").cpu()
+net_c = PretrainedConvNext("convnext_small_in22k").cuda()
 net_c.load_state_dict(torch.load('./classifier_32.pt')['icnn'])
 net_i.eval()
 net_c.eval()
@@ -25,8 +25,8 @@ def align(x1):
 def predict(img):
     with torch.no_grad():
         img=align(img)
-        image_tensor=TF.to_tensor(img).cpu()
-        image_tensor=image_tensor.unsqueeze(0).cpu()
+        image_tensor=TF.to_tensor(img).cuda()
+        image_tensor=image_tensor.unsqueeze(0).cuda()
         ipt=net_c(image_tensor)
         image_tensor = image_tensor.half()
         ipt = ipt.half()