Fix device usage?

app.py

@@ -10,7 +10,6 @@ from align.detector import detect_faces
 from align.align_trans import get_reference_facial_points, warp_and_crop_face
 import torchvision.transforms as transforms
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 to_tensor = transforms.ToTensor()
 
 eps = 0.05
@@ -45,6 +44,8 @@ for root in model_roots:
 
 @spaces.GPU(duration=120)
 def protect(img, progress=gr.Progress(track_tqdm=True)):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
     models_attack, V_reduction, dim = prepare_models(
         model_backbones,
         input_size,

util/attack_utils.py

@@ -55,6 +55,7 @@ class Attack(nn.Module):
 
     def execute(self, images, dir_vec, direction):
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        print("Device in Excute:", device)
         self.lpips_loss.to(device)
         images = Variable(images).to(device)
         dir_vec = dir_vec.to(device)

util/feature_extraction_utils.py

@@ -5,11 +5,10 @@ import torchvision.transforms as transforms
 import torch.nn as nn
 import numpy as np
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
 
 def warp_image(tensor_img, theta_warp, crop_size=112):
     # applies affine transform theta to image and crops it
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     theta_warp = torch.Tensor(theta_warp).unsqueeze(0).to(device)
     grid = F.affine_grid(theta_warp, tensor_img.size())

util/prepare_utils.py

@@ -123,6 +123,8 @@ def get_ensemble(
     warp=False,
     theta_warp=None,
 ):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
     # function prepares ensemble of feature extractors
     # outputs list of pytorch nn models
     feature_extractor_ensemble = []
@@ -213,6 +215,8 @@ def prepare_models(
     using_subspace,
     V_reduction_root,
 ):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
     backbone_dict = {
         "IR_50": IR_50(input_size),
         "IR_152": IR_152(input_size),