edit util.py and gradio_demo.py

FONT/modules/util.py

@@ -306,9 +306,9 @@ class AT_net(nn.Module):
 
 
 class AT_net2(nn.Module):
-    def __init__(self):
+    def __init__(self, device):
         super(AT_net2, self).__init__()
-
+        self.device = device
         down_blocks = []
         for i in range(8):
             down_blocks.append(DownBlock2d(3 if i == 0 else  2 * (2 ** i),
@@ -372,8 +372,16 @@ class AT_net2(nn.Module):
 
 
     def forward(self, example_image, audio, pose, jaco_net, weight):
-        hidden = ( torch.autograd.Variable(torch.zeros(3, audio.size(0), 256).cuda()),
-                      torch.autograd.Variable(torch.zeros(3, audio.size(0), 256).cuda()))
+
+        hidden_ele1 = torch.zeros(3, audio.size(0), 256)
+        hidden_ele2 = torch.zeros(3, audio.size(0), 256)
+
+        if self.device == 'cuda':
+            hidden_ele1 = hidden_ele1.cuda()
+            hidden_ele2 = hidden_ele2.cuda()
+
+
+        hidden = (torch.autograd.Variable(hidden_ele1), torch.autograd.Variable(hidden_ele2) )
         outs = example_image
         for down_block in self.down_blocks:
             outs = down_block(outs)