Update e4e_projection.py

e4e_projection.py

@@ -8,15 +8,17 @@ 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(
         [