fix cuda -> cpu

Embed_Secret.py

@@ -130,7 +130,7 @@ def embed_secret(model_name, model, cover, tform, secret):
     if model_name == 'UNet':
         w, h = cover.size
         with torch.no_grad():
-            im = tform(cover).unsqueeze(0).cuda()  # 1, 3, 256, 256
+            im = tform(cover).unsqueeze(0).to(model.device)  # 1, 3, 256, 256
             stego, _ = model(im, secret)  # 1, 3, 256, 256
             res = (stego.clamp(-1,1) - im)  # (1,3,256,256) residual
             res = torch.nn.functional.interpolate(res, (h,w), mode='bilinear')
@@ -147,7 +147,7 @@ def identity(x):
 def decode_secret(model_name, model, im, tform):
     if model_name in ['RoSteALS', 'UNet']:
         with torch.no_grad():
-            im = tform(im).unsqueeze(0).cuda()  # 1, 3, 256, 256
+            im = tform(im).unsqueeze(0).to(model.device)  # 1, 3, 256, 256
             secret_pred = (model.decoder(im) > 0).cpu().numpy()  # 1, 100
     else:
         raise NotImplementedError
@@ -241,7 +241,7 @@ def app(args):
     ])
     if image_file is not None and secret_text is not None:
         secret = ecc.encode_text([secret_text])  # (1, len)
-        secret = torch.from_numpy(secret).float().cuda()
+        secret = torch.from_numpy(secret).float().to(model.device)
         # im = tform(im).unsqueeze(0).cuda()  # (1,3,H,W)
         stego = embed_secret(model_name, model, im, tform_emb, secret)
         st.image(stego, width=display_width)