update

app.py

@@ -57,16 +57,16 @@ def initialize_models():
                                                 rep_drop_prob=0.,
                                                 use_class_label=False)
     generator = generator.to(device)
-    checkpoint = torch.load(generator_weights, map_location='cpu')
+    checkpoint = torch.load(generator_weights, map_location=device)
     generator.load_state_dict(checkpoint['model_vec2face'])
     generator.eval()
 
     id_model = iresnet("100", fp16=True).to(device)
-    id_model.load_state_dict(torch.load(id_model_weights, map_location='cpu'))
+    id_model.load_state_dict(torch.load(id_model_weights, map_location=device))
     id_model.eval()
 
     quality_model = iresnet("100", fp16=True).to(device)
-    quality_model.load_state_dict(torch.load(quality_model_weights, map_location='cpu'))
+    quality_model.load_state_dict(torch.load(quality_model_weights, map_location=device))
     quality_model.eval()
 
     pose_model = SixDRepNet(backbone_name='RepVGG-B1g2',