Update e4e_projection.py

e4e_projection.py

@@ -8,17 +8,18 @@ from argparse import Namespace
 from e4e.models.psp import pSp
 from util import *
 
-device='cpu'
-model_path = 'models/e4e_ffhq_encode.pt'
-ckpt = torch.load(model_path, map_location='cpu')
-opts = ckpt['opts']
-opts['checkpoint_path'] = model_path
-opts= Namespace(**opts)
-net = pSp(opts, device).eval().to(device)
+
 
 @ torch.no_grad()
 def projection(img, name, device='cuda'):
 
+    model_path = 'models/e4e_ffhq_encode.pt'
+    ckpt = torch.load(model_path, map_location='cpu')
+    opts = ckpt['opts']
+    opts['checkpoint_path'] = model_path
+    opts= Namespace(**opts)
+    net = pSp(opts, device).eval().to(device)
+
 
     transform = transforms.Compose(
         [