udpate: app.py import FSGenerator

inference/tricks.py

@@ -138,10 +138,14 @@ class SoftErosion(nn.Module):
         return x, mask
 
 
+if torch.cuda.is_available():
+    device = torch.device(0)
+else:
+    device = torch.device('cpu')
 vgg_mean = torch.tensor([[[0.485]], [[0.456]], [[0.406]]],
-                             requires_grad=False, device=torch.device(0))
+                             requires_grad=False, device=device)
 vgg_std = torch.tensor([[[0.229]], [[0.224]], [[0.225]]],
-                            requires_grad=False, device=torch.device(0))
+                            requires_grad=False, device=device)
 def load_bisenet():
     bisenet_model = BiSeNet(n_classes=19)
     bisenet_model.load_state_dict(